Refactor TEDLIUM unittest (#1135)

test/torchaudio_unittest/datasets/tedlium_test.py

 import os
 import platform
-import unittest
 from pathlib import Path
 
-from torchaudio.datasets import tedlium
-
 from torchaudio_unittest.common_utils import (
-    TestBaseMixin,
     TempDirMixin,
     TorchaudioTestCase,
     get_whitenoise,
     save_wav,
-    normalize_wav,
 )
 
+from torchaudio.datasets import tedlium
+
 # Used to generate a unique utterance for each dummy audio file
-UTTERANCES = [
+_UTTERANCES = [
     "AaronHuey_2010X 1 AaronHuey_2010X 0.0 2.0 <o,f0,female> script1\n",
     "AaronHuey_2010X 1 AaronHuey_2010X 2.0 4.0 <o,f0,female> script2\n",
     "AaronHuey_2010X 1 AaronHuey_2010X 4.0 6.0 <o,f0,female> script3\n",
@@ -23,7 +20,7 @@ UTTERANCES = [
     "AaronHuey_2010X 1 AaronHuey_2010X 8.0 10.0 <o,f0,female> script5\n",
 ]
 
-PHONEME = [
+_PHONEME = [
     "a AH",
     "a(2) EY",
     "aachen AA K AH N",
@@ -34,6 +31,65 @@ PHONEME = [
 ]
 
 
+def get_mock_dataset(dataset_dir):
+    """
+    dataset_dir: directory of the mocked dataset
+    """
+    mocked_samples = {}
+    os.makedirs(dataset_dir, exist_ok=True)
+    sample_rate = 16000  # 16kHz
+    seed = 0
+
+    for release in ["release1", "release2", "release3"]:
+        data = get_whitenoise(sample_rate=sample_rate, duration=10.00, n_channels=1, dtype="float32", seed=seed)
+        if release in ["release1", "release2"]:
+            release_dir = os.path.join(
+                dataset_dir,
+                tedlium._RELEASE_CONFIGS[release]["folder_in_archive"],
+                tedlium._RELEASE_CONFIGS[release]["subset"],
+            )
+        else:
+            release_dir = os.path.join(
+                dataset_dir,
+                tedlium._RELEASE_CONFIGS[release]["folder_in_archive"],
+                tedlium._RELEASE_CONFIGS[release]["data_path"],
+            )
+        os.makedirs(release_dir, exist_ok=True)
+        os.makedirs(os.path.join(release_dir, "stm"), exist_ok=True)  # Subfolder for transcripts
+        os.makedirs(os.path.join(release_dir, "sph"), exist_ok=True)  # Subfolder for audio files
+        filename = f"{release}.sph"
+        path = os.path.join(os.path.join(release_dir, "sph"), filename)
+        save_wav(path, data, sample_rate)
+
+        trans_filename = f"{release}.stm"
+        trans_path = os.path.join(os.path.join(release_dir, "stm"), trans_filename)
+        with open(trans_path, "w") as f:
+            f.write("".join(_UTTERANCES))
+
+        dict_filename = f"{release}.dic"
+        dict_path = os.path.join(release_dir, dict_filename)
+        with open(dict_path, "w") as f:
+            f.write("\n".join(_PHONEME))
+
+        # Create a samples list to compare with
+        mocked_samples[release] = []
+        for utterance in _UTTERANCES:
+            talk_id, _, speaker_id, start_time, end_time, identifier, transcript = utterance.split(" ", 6)
+            start_time = int(float(start_time)) * sample_rate
+            end_time = int(float(end_time)) * sample_rate
+            sample = (
+                data[:, start_time:end_time],
+                sample_rate,
+                transcript,
+                talk_id,
+                speaker_id,
+                identifier,
+            )
+            mocked_samples[release].append(sample)
+        seed += 1
+    return mocked_samples
+
+
 class Tedlium(TempDirMixin):
     root_dir = None
     samples = {}
@@ -42,57 +98,7 @@ class Tedlium(TempDirMixin):
     def setUpClass(cls):
         cls.root_dir = cls.get_base_temp_dir()
         cls.root_dir = dataset_dir = os.path.join(cls.root_dir, "tedlium")
-        os.makedirs(dataset_dir, exist_ok=True)
-        sample_rate = 16000  # 16kHz
-        seed = 0
-
-        for release in ["release1", "release2", "release3"]:
-            data = get_whitenoise(sample_rate=sample_rate, duration=10.00, n_channels=1, dtype="float32", seed=seed)
-            if release in ["release1", "release2"]:
-                release_dir = os.path.join(
-                    dataset_dir,
-                    tedlium._RELEASE_CONFIGS[release]["folder_in_archive"],
-                    tedlium._RELEASE_CONFIGS[release]["subset"],
-                )
-            else:
-                release_dir = os.path.join(
-                    dataset_dir,
-                    tedlium._RELEASE_CONFIGS[release]["folder_in_archive"],
-                    tedlium._RELEASE_CONFIGS[release]["data_path"],
-                )
-            os.makedirs(release_dir, exist_ok=True)
-            os.makedirs(os.path.join(release_dir, "stm"), exist_ok=True)  # Subfolder for transcripts
-            os.makedirs(os.path.join(release_dir, "sph"), exist_ok=True)  # Subfolder for audio files
-            filename = f"{release}.sph"
-            path = os.path.join(os.path.join(release_dir, "sph"), filename)
-            save_wav(path, data, sample_rate)
-
-            trans_filename = f"{release}.stm"
-            trans_path = os.path.join(os.path.join(release_dir, "stm"), trans_filename)
-            with open(trans_path, "w") as f:
-                f.write("".join(UTTERANCES))
-
-            dict_filename = f"{release}.dic"
-            dict_path = os.path.join(release_dir, dict_filename)
-            with open(dict_path, "w") as f:
-                f.write("\n".join(PHONEME))
-
-            # Create a samples list to compare with
-            cls.samples[release] = []
-            for utterance in UTTERANCES:
-                talk_id, _, speaker_id, start_time, end_time, identifier, transcript = utterance.split(" ", 6)
-                start_time = int(float(start_time)) * sample_rate
-                end_time = int(float(end_time)) * sample_rate
-                sample = (
-                    data[:, start_time:end_time],
-                    sample_rate,
-                    transcript,
-                    talk_id,
-                    speaker_id,
-                    identifier,
-                )
-                cls.samples[release].append(sample)
-            seed += 1
+        cls.samples = get_mock_dataset(dataset_dir)
 
     def _test_tedlium(self, dataset, release):
         num_samples = 0
@@ -110,7 +116,7 @@ class Tedlium(TempDirMixin):
         dataset._dict_path = os.path.join(dataset._path, f"{release}.dic")
         phoneme_dict = dataset.phoneme_dict
         phoenemes = [f"{key} {' '.join(value)}" for key, value in phoneme_dict.items()]
-        assert phoenemes == PHONEME
+        assert phoenemes == _PHONEME
 
     def test_tedlium_release1_str(self):
         release = "release1"