Fix unit tests

Author: schustmi

tests/unit/config/test_compiler.py

@@ -34,8 +34,7 @@ def test_compiling_pipeline_with_invalid_run_name_fails(
 ):
     """Tests that compiling a pipeline with an invalid run name fails."""
     pipeline_instance = empty_pipeline
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     with pytest.raises(ValueError):
         Compiler().compile(
             pipeline=pipeline_instance,
@@ -54,8 +53,7 @@ def _no_step_pipeline():
 def test_compiling_pipeline_without_steps_fails(local_stack):
     """Tests that compiling a pipeline without steps fails."""
     pipeline_instance = _no_step_pipeline
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     with pytest.raises(ValueError):
         Compiler().compile(
             pipeline=pipeline_instance,
@@ -71,8 +69,7 @@ def test_compiling_pipeline_with_missing_step_operator(
     pipeline_instance = one_step_pipeline(
         empty_step.configure(step_operator="s")
     )
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     with pytest.raises(StackValidationError):
         Compiler().compile(
             pipeline=pipeline_instance,
@@ -89,8 +86,7 @@ def test_compiling_pipeline_with_missing_experiment_tracker(
     pipeline_instance = one_step_pipeline(
         empty_step.configure(experiment_tracker="e")
     )
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     with pytest.raises(StackValidationError):
         Compiler().compile(
             pipeline=pipeline_instance,
@@ -113,8 +109,7 @@ def test_pipeline_and_steps_dont_get_modified_during_compilation(
             "_empty_step": StepConfigurationUpdate(extra={"key": "new_value"})
         },
     )
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     Compiler().compile(
         pipeline=pipeline_instance,
         stack=local_stack,
@@ -148,8 +143,7 @@ def pipeline_instance():
         empty_step(id="step_1")
         empty_step(id="step_2", after="step_1")
 
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     snapshot = Compiler().compile(
         pipeline=pipeline_instance,
         stack=local_stack,
@@ -204,8 +198,7 @@ class StubSettings(BaseSettings):
             )
         },
     )
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     snapshot = Compiler().compile(
         pipeline=pipeline_instance,
         stack=local_stack,
@@ -252,8 +245,7 @@ def test_general_settings_merging(one_step_pipeline, empty_step, local_stack):
             )
         },
     )
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     snapshot = Compiler().compile(
         pipeline=pipeline_instance,
         stack=local_stack,
@@ -299,8 +291,7 @@ def test_extra_merging(one_step_pipeline, empty_step, local_stack):
         steps={"_empty_step": StepConfigurationUpdate(extra=run_step_extra)},
     )
 
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
 
     snapshot = Compiler().compile(
         pipeline=pipeline_instance,
@@ -358,8 +349,7 @@ def test_success_hook_merging(
         },
     )
 
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     snapshot = Compiler().compile(
         pipeline=pipeline_instance,
         stack=local_stack,
@@ -409,8 +399,7 @@ def test_failure_hook_merging(
         },
     )
 
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     snapshot = Compiler().compile(
         pipeline=pipeline_instance,
         stack=local_stack,
@@ -453,8 +442,7 @@ def test_stack_component_settings_for_missing_component_are_ignored(
         steps={"_empty_step": StepConfigurationUpdate(settings=settings)},
     )
 
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     snapshot = Compiler().compile(
         pipeline=pipeline_instance,
         stack=local_stack,
@@ -500,8 +488,7 @@ class StubSettings(BaseSettings):
         steps={"_empty_step": StepConfigurationUpdate(settings=settings)},
     )
 
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     snapshot = Compiler().compile(
         pipeline=pipeline_instance,
         stack=local_stack,
@@ -551,8 +538,7 @@ class StubSettings(BaseSettings):
         steps={"_empty_step": StepConfigurationUpdate(settings=settings)},
     )
 
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     snapshot = Compiler().compile(
         pipeline=pipeline_instance,
         stack=local_stack,
@@ -585,8 +571,7 @@ def test_spec_compilation(local_stack):
     def pipeline_instance():
         s2(s1())
 
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
     spec = (
         Compiler()
         .compile(
@@ -652,8 +637,7 @@ class StubSettings(BaseSettings):
         settings={"orchestrator": shortcut_settings}
     )
 
-    with pipeline_instance_with_shortcut_settings:
-        pipeline_instance_with_shortcut_settings.entrypoint()
+    pipeline_instance_with_shortcut_settings.prepare()
 
     with does_not_raise():
         snapshot = Compiler().compile(
@@ -677,8 +661,7 @@ class StubSettings(BaseSettings):
         }
     )
 
-    with pipeline_instance_with_duplicate_settings:
-        pipeline_instance_with_duplicate_settings.entrypoint()
+    pipeline_instance_with_duplicate_settings.prepare()
 
     with pytest.raises(ValueError):
         snapshot = Compiler().compile(

tests/unit/conftest.py

@@ -294,7 +294,6 @@ def step_context_with_no_output(
     sample_pipeline_run: PipelineRunResponse,
     sample_step_run: StepRunResponse,
 ) -> StepContext:
-    StepContext._clear()
     return StepContext(
         pipeline_run=sample_pipeline_run,
         step_run=sample_step_run,
@@ -312,7 +311,6 @@ def step_context_with_single_output(
     materializers = {"output_1": (BaseMaterializer,)}
     artifact_uris = {"output_1": ""}
     artifact_configs = {"output_1": None}
-    StepContext._clear()
     return StepContext(
         pipeline_run=sample_pipeline_run,
         step_run=sample_step_run,
@@ -337,7 +335,6 @@ def step_context_with_two_outputs(
     }
     artifact_configs = {"output_1": None, "output_2": None}
 
-    StepContext._clear()
     return StepContext(
         pipeline_run=sample_pipeline_run,
         step_run=sample_step_run,
@@ -634,6 +631,7 @@ def sample_snapshot_response_model(
             updated=datetime.now(),
             runnable=True,
             deployable=True,
+            is_dynamic=False,
         ),
         metadata=PipelineSnapshotResponseMetadata(
             run_name_template="",

tests/unit/entrypoints/test_base_entrypoint_configuration.py

@@ -69,4 +69,4 @@ def test_loading_the_snapshot(clean_client):
         arguments=["--snapshot_id", str(snapshot.id)]
     )
 
-    assert entrypoint_config.load_snapshot() == snapshot
+    assert entrypoint_config.snapshot == snapshot

tests/unit/orchestrators/test_publish_utils.py

@@ -100,6 +100,7 @@ def test_pipeline_run_status_computation(
             run_status=ExecutionStatus.RUNNING,
             step_statuses=step_statuses,
             num_steps=num_steps,
+            is_dynamic_pipeline=False,
         )
         == expected_run_status
     )

tests/unit/orchestrators/test_step_runner.py

@@ -22,7 +22,11 @@
 from zenml.config.step_configurations import Step
 from zenml.config.step_run_info import StepRunInfo
 from zenml.enums import ArtifactSaveType
-from zenml.models import PipelineRunResponse, StepRunResponse
+from zenml.models import (
+    PipelineRunResponse,
+    PipelineSnapshotResponse,
+    StepRunResponse,
+)
 from zenml.orchestrators.step_launcher import StepRunner
 from zenml.stack import Stack
 from zenml.steps import step
@@ -43,6 +47,7 @@ def test_running_a_successful_step(
     local_stack,
     sample_pipeline_run: PipelineRunResponse,
     sample_step_run: StepRunResponse,
+    sample_snapshot_response_model: PipelineSnapshotResponse,
 ):
     """Tests that running a successful step runs the step entrypoint
     and correctly prepares/cleans up."""
@@ -74,7 +79,9 @@ def test_running_a_successful_step(
         run_name="run_name",
         pipeline_step_name="step_name",
         config=step.config,
+        spec=step.spec,
         pipeline=pipeline_config,
+        snapshot=sample_snapshot_response_model,
         force_write_logs=lambda: None,
     )
 
@@ -98,6 +105,7 @@ def test_running_a_failing_step(
     local_stack,
     sample_pipeline_run: PipelineRunResponse,
     sample_step_run: StepRunResponse,
+    sample_snapshot_response_model: PipelineSnapshotResponse,
 ):
     """Tests that running a failing step runs the step entrypoint
     and correctly prepares/cleans up."""
@@ -130,7 +138,9 @@ def test_running_a_failing_step(
         run_name="run_name",
         pipeline_step_name="step_name",
         config=step.config,
+        spec=step.spec,
         pipeline=pipeline_config,
+        snapshot=sample_snapshot_response_model,
         force_write_logs=lambda: None,
     )
 

tests/unit/pipelines/test_base_pipeline.py

@@ -429,8 +429,7 @@ def test_compiling_a_pipeline_merges_schedule(
     config_path.write_text(run_config.yaml())
 
     pipeline_instance = empty_pipeline
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
 
     _, schedule, _ = pipeline_instance._compile(
         config_path=str(config_path),
@@ -462,8 +461,7 @@ def test_compiling_a_pipeline_merges_build(
     in_code_build_id = uuid4()
 
     pipeline_instance = empty_pipeline
-    with pipeline_instance:
-        pipeline_instance.entrypoint()
+    pipeline_instance.prepare()
 
     # Config with ID
     _, _, build = pipeline_instance._compile(

tests/unit/stack/test_stack.py

@@ -149,8 +149,7 @@ def test_stack_submission(
     components."""
     # Mock the pipeline run registering which tries (and fails) to serialize
     # our mock objects
-    with empty_pipeline:
-        empty_pipeline.entrypoint()
+    empty_pipeline.prepare()
     snapshot = Compiler().compile(
         pipeline=empty_pipeline,
         stack=stack_with_mock_components,

tests/unit/steps/test_step_context.py

@@ -22,22 +22,15 @@
 from zenml.steps import StepContext
 
 
-def test_step_context_is_singleton(step_context_with_no_output):
-    """Tests that the step context is a singleton."""
-    assert StepContext() is step_context_with_no_output
-
-
 def test_get_step_context(step_context_with_no_output):
     """Unit test for `get_step_context()`."""
 
-    # step context exists -> returns the step context
-    assert get_step_context() is StepContext()
-
-    # step context does not exist -> raises an exception
-    StepContext._clear()
     with pytest.raises(RuntimeError):
         get_step_context()
 
+    with step_context_with_no_output:
+        assert get_step_context() is step_context_with_no_output
+
 
 def test_initialize_step_context_with_mismatched_keys(
     sample_pipeline_run,
@@ -49,7 +42,6 @@ def test_initialize_step_context_with_mismatched_keys(
     artifact_configs = {"some_yet_another_output_name": None}
 
     with pytest.raises(StepContextError):
-        StepContext._clear()
         StepContext(
             pipeline_run=sample_pipeline_run,
             step_run=sample_step_run,
@@ -69,7 +61,6 @@ def test_initialize_step_context_with_matching_keys(
     artifact_configs = {"some_output_name": None}
 
     with does_not_raise():
-        StepContext._clear()
         StepContext(
             pipeline_run=sample_pipeline_run,
             step_run=sample_step_run,