Refactor dataset preparation fixture to avoid redundancy and limit test parametrization to reduce time (#632)

* refactor combined_dataset fixture for reusable dataset preparation

* remove repeated tests

Author: bhimrazy

tests/conftest.py

@@ -52,19 +52,24 @@ def mosaic_mds_index_data():
     }
 
 
-@pytest.fixture
-def combined_dataset(tmpdir_factory):
-    tmpdir = tmpdir_factory.mktemp("data")
+@pytest.fixture(scope="session")
+def prepare_combined_dataset(tmpdir_factory):
+    tmpdir = tmpdir_factory.mktemp("combined_dataset")
     datasets = [str(tmpdir.join(f"dataset_{i}")) for i in range(2)]
     for dataset in datasets:
         cache = Cache(input_dir=dataset, chunk_bytes="64MB")
         for i in range(50):
             cache[i] = i
         cache.done()
         cache.merge()
+    return datasets
 
-    dataset_1 = StreamingDataset(datasets[0], shuffle=True)
-    dataset_2 = StreamingDataset(datasets[1], shuffle=True)
+
+@pytest.fixture
+def combined_dataset(prepare_combined_dataset):
+    dataset_1_path, dataset_2_path = prepare_combined_dataset
+    dataset_1 = StreamingDataset(dataset_1_path)
+    dataset_2 = StreamingDataset(dataset_2_path)
     return CombinedStreamingDataset(datasets=[dataset_1, dataset_2])
 
 

tests/streaming/test_combined.py

@@ -537,7 +537,6 @@ def test_combined_dataset_dataloader_states_without_any_iterations(combined_data
 @pytest.mark.timeout(120)
 @pytest.mark.parametrize("num_workers", [0, 2, 4])
 def test_combined_dataset_dataloader_states_complete_iterations(combined_dataset, num_workers):
-    print(f"Testing with num_workers={num_workers}")
     dataloader = StreamingDataLoader(combined_dataset, batch_size=4, num_workers=num_workers)
     assert len(dataloader) == 25, "Dataloader length should be 25 (50+50 items / batch size 4)"
 
@@ -559,16 +558,13 @@ def test_combined_dataset_dataloader_states_complete_iterations(combined_dataset
 
 
 @pytest.mark.timeout(300)
-@pytest.mark.parametrize(("num_workers", "break_at"), [(0, 10), (0, 15), (2, 10), (2, 15), (4, 10), (4, 15)])
+@pytest.mark.parametrize(("num_workers", "break_at"), [(0, 10), (0, 15), (2, 15), (4, 15)])
 def test_combined_dataset_dataloader_states_partial_iterations(combined_dataset, num_workers, break_at):
-    print(f"Testing with num_workers={num_workers}, break_at={break_at}")
-
     # Verify dataloader state after partial last iteration
     dataloader = StreamingDataLoader(combined_dataset, batch_size=4, num_workers=num_workers)
 
     total_batches = len(dataloader)
     assert total_batches == 25, "Dataloader length should be 25 (100 items / batch size 4)"
-
     assert not dataloader.restore, "Dataloader should not be in restore state initially."
 
     # Partial iteration up to 'break_at'