update output df

epoch8 · Aug 22, 2024 · 6f55b3c · 6f55b3c
1 parent c91a097
commit 6f55b3c
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/tests/test_complex_pipeline.py b/tests/test_complex_pipeline.py
@@ -409,12 +409,11 @@ def get_some_prediction_only_on_best_model(
     )
     steps = build_compute(ds, catalog, pipeline)
     run_steps(ds, steps)
-    test__df_output = get_some_prediction_only_on_best_model(
-        df__image=test_df__image,
-        df__image__attribute=test_df__image__attribute,
-        df__prediction=test_df__prediction,
-        df__best_model=test_df__best_model
-    )
+    test__df_output = pd.DataFrame({
+        "image_id": range(N),
+        "model_id": [4] * N,
+        "result": [0] * N
+    })
     assert_df_equal(
         ds.get_table("tbl_output").get_data(),
         test__df_output,