Fixing unit tests

astronomer · Dec 17, 2024 · cb6ad85 · cb6ad85
1 parent a1d15a2
commit cb6ad85
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/tests/test_dagfactory.py b/tests/test_dagfactory.py
@@ -441,10 +441,9 @@ def test_generate_dags_with_removal_valid_and_callback():
 
 
 def test_set_callback_after_loading_config():
-    td = dagfactory.DagFactory(config=DAG_FACTORY_CONFIG)
-    td.config["default"]["default_args"]["on_success_callback"] = dagfactory.DagFactory(
-        config=DAG_FACTORY_CONFIG
-    ).build_dags
+    td = dagfactory.DagFactory(config=DAG_FACTORY_CONFIG)  # Generate the DAG factory object
+    td.config["default"]["default_args"]["on_success_callback"] = f"{__name__}.print_context_callback"
+    td.clean_dags(globals())
     td.generate_dags(globals())