Wrapped test in flow

MatsMoll · May 20, 2024 · 7a536e5 · 7a536e5
1 parent 0d7d033
commit 7a536e5
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tests/test_train_pipeline.py b/tests/test_train_pipeline.py
@@ -4,6 +4,7 @@
 from aligned.data_source.batch_data_source import DummyDataSource
 from sklearn.ensemble import RandomForestClassifier
 
+from prefect import flow
 from prefect.testing.utilities import prefect_test_harness
 
 from src.model_registry import InMemoryModelRegristry
@@ -36,7 +37,8 @@ async def test_generic_classifier_train_pipeline_using_prefect():
     tracker = StdoutExperimentTracker()
     model = RandomForestClassifier()
 
-    with prefect_test_harness():
+    @flow(name="test")
+    def test_flow():
         await classifier_from_train_test_validation_set(
             store=store,
             model_contract=MovieReviewIsNegative.metadata.name,
@@ -52,5 +54,8 @@ async def test_generic_classifier_train_pipeline_using_prefect():
             tracker=tracker,
         )
 
+    with prefect_test_harness():
+        test_flow()
+
     assert len(registry.models) == 1