fix pytest

VIDA-NYU · Jun 24, 2024 · 6c7c4e6 · 6c7c4e6
1 parent b6a20ba
commit 6c7c4e6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -29,13 +29,13 @@ def test_sample_dataset():
     dataset = pd.read_csv(dataset_path)
     X = dataset.drop(columns=["rating"])
     y = dataset[["rating"]]
-    sample_size = 10
+    sample_size = 100 # 10 * 10
 
     actual_X, actual_y, actual_is_sampled = sample_dataset(
         X, y, sample_size, "CLASSIFICATION"
     )
-    expected_X_len = sample_size
-    expected_y_len = sample_size
+    expected_X_len = sample_size // len(X.columns)
+    expected_y_len = sample_size // len(X.columns)
     expected_is_sampled = True
 
     assert actual_is_sampled == expected_is_sampled