diff --git a/test/datasets/test_stsb.py b/test/datasets/test_stsb.py
index 7e95d4f5ca..f74cd2b6d2 100644
--- a/test/datasets/test_stsb.py
+++ b/test/datasets/test_stsb.py
@@ -21,7 +21,7 @@ def _get_mock_dataset(root_dir):
 
     seed = 1
     mocked_data = defaultdict(list)
-    for file_name, name in zip(["sts-train.csv", "sts-dev.csv" "sts-test.csv"], ["train", "dev", "test"]):
+    for file_name, name in zip(["sts-train.csv", "sts-dev.csv", "sts-test.csv"], ["train", "dev", "test"]):
         txt_file = os.path.join(temp_dataset_dir, file_name)
         with open(txt_file, "w", encoding="utf-8") as f:
             for i in range(5):
diff --git a/test/datasets/test_wikitexts.py b/test/datasets/test_wikitexts.py
index ee8bfafb8f..4706a53093 100644
--- a/test/datasets/test_wikitexts.py
+++ b/test/datasets/test_wikitexts.py
@@ -25,12 +25,12 @@ def _get_mock_dataset(root_dir, base_dir_name):
     file_names = ("wiki.train.tokens", "wiki.valid.tokens", "wiki.test.tokens")
     for file_name in file_names:
         csv_file = os.path.join(temp_dataset_dir, file_name)
-        mocked_lines = mocked_data[os.path.splitext(file_name)[0]]
+        mocked_lines = mocked_data[file_name.split(".")[1]]
         with open(csv_file, "w", encoding="utf-8") as f:
             for i in range(5):
                 rand_string = get_random_unicode(seed)
-                dataset_line = rand_string
-                f.write(f"{rand_string}\n")
+                dataset_line = f"{rand_string}\n"
+                f.write(dataset_line)
 
                 # append line to correct dataset split
                 mocked_lines.append(dataset_line)
@@ -38,15 +38,17 @@ def _get_mock_dataset(root_dir, base_dir_name):
 
     if base_dir_name == WikiText103.__name__:
         compressed_file = "wikitext-103-v1"
+        arcname_folder = "wikitext-103"
     else:
         compressed_file = "wikitext-2-v1"
+        arcname_folder = "wikitext-2"
 
     compressed_dataset_path = os.path.join(base_dir, compressed_file + ".zip")
     # create zip file from dataset folder
     with zipfile.ZipFile(compressed_dataset_path, "w") as zip_file:
         for file_name in file_names:
             txt_file = os.path.join(temp_dataset_dir, file_name)
-            zip_file.write(txt_file, arcname=compressed_file)
+            zip_file.write(txt_file, arcname=os.path.join(arcname_folder, file_name))
 
     return mocked_data