lint test, fix bug in test where TrackingDataset.prepare isn't called

talmolab · Aug 30, 2024 · f8fdcec · f8fdcec
1 parent 0677ccd
commit f8fdcec
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
@@ -534,6 +534,7 @@ def test_train_test_split(two_flies):
 
     splits = (0.7, 0.2, 0.1)
     tracking_ds = TrackingDataset(train_ds=train_ds, splits=splits)
+    tracking_ds.setup("fit")
 
     assert len(tracking_ds.train_ds) == int(ds_length * splits[0])
     assert len(tracking_ds.val_ds) == int(ds_length * splits[1])
@@ -544,6 +545,8 @@ def test_train_test_split(two_flies):
     assert tracking_ds.test_ds.dataset.augmentations is None
 
     tracking_ds = TrackingDataset(train_ds=train_ds, splits=splits[:-1])
+    tracking_ds.setup("fit")
+
     assert len(tracking_ds.train_ds) == int(ds_length * splits[0])
     assert len(tracking_ds.val_ds) == int(ds_length * splits[1])
     assert tracking_ds.test_ds is None
@@ -558,12 +561,14 @@ def test_train_test_split(two_flies):
     ds_length = len(val_ds)
 
     splits = (0.5, 0.5)
+
     tracking_ds = TrackingDataset(train_ds=train_ds, val_ds=val_ds, splits=splits)
+    tracking_ds.setup("fit")
 
     assert len(tracking_ds.train_ds) == len(train_ds)
     assert len(tracking_ds.val_ds) == int(ds_length * splits[0])
     assert len(tracking_ds.test_ds) == int(ds_length * splits[1])
 
     assert tracking_ds.train_ds.augmentations is not None
     assert tracking_ds.val_ds.dataset.augmentations is None
-    assert tracking_ds.test_ds.dataset.augmentations is None
+    assert tracking_ds.test_ds.dataset.augmentations is None