Fix filename tests

Signed-off-by: Ryan Wolf <[email protected]>
NVIDIA · Apr 23, 2024 · 89009a8 · 89009a8
1 parent 491f8b6
commit 89009a8
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/tests/test_shuffle.py b/tests/test_shuffle.py
@@ -53,11 +53,11 @@ def test_filename(self):
         expected_data = {
             "text": ["one", "three", "four", "five", "two"],
             "filename": [
+                "file_0000000000.jsonl",
+                "file_0000000000.jsonl",
+                "file_0000000000.jsonl",
                 "file_0000000001.jsonl",
                 "file_0000000001.jsonl",
-                "file_0000000001.jsonl",
-                "file_0000000002.jsonl",
-                "file_0000000002.jsonl",
             ],
         }
         pdf = pd.DataFrame(expected_data)
@@ -76,11 +76,11 @@ def test_custom_filenames(self):
         expected_data = {
             "text": ["one", "three", "four", "five", "two"],
             "filename": [
+                "my_0.test",
+                "my_0.test",
+                "my_0.test",
                 "my_1.test",
                 "my_1.test",
-                "my_1.test",
-                "my_2.test",
-                "my_2.test",
             ],
         }
         pdf = pd.DataFrame(expected_data)