fix: flaky test

MatsMoll · Nov 6, 2023 · 6a5fd2b · 6a5fd2b
1 parent ed69ece
commit 6a5fd2b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/aligned/tests/test_model_target.py b/aligned/tests/test_model_target.py
@@ -124,7 +124,7 @@ class TestModel:
 
     await store.insert_into(FeatureLocation.model('test_model'), {'id': [1, 2, 3], 'a': [10, 14, 20]})
 
-    stored_data = pl.read_parquet(path)
+    stored_data = pl.read_parquet(path).select(expected_frame.columns)
     assert stored_data.frame_equal(expected_frame)