Fix collecting absent imports in pytest (#1154)

.github/actions/test-e2e/action.yml

@@ -30,7 +30,8 @@ runs:
 
     - name: E2E - ${{ inputs.module }}
       run: |
-        coverage run -p -m pytest -v ./tests/e2e/ -m "${{ inputs.module }}" \
+        coverage run -p -m \
+        pytest -v ./tests/e2e/ -m "${{ inputs.module }}" \
         --timeout=600 --timeout_method=thread \
         --junitxml="./test-results/test-e2e-${{ inputs.name }}-${{ inputs.module }}-${{ inputs.os }}-${{ inputs.python-version }}.xml"
       shell: bash

tests/e2e/integrations/test_zenml.py

@@ -1,62 +1,54 @@
 import hashlib
 
-from pytest import (
-    fixture,
-    mark,
-)
+import pytest
 from sklearn.datasets import load_digits
 from sklearn.model_selection import train_test_split
 from sklearn.svm import SVC
-from zenml.client import Client
-from zenml.enums import StackComponentType
-from zenml.exceptions import (
-    InitializationException,
-    StackComponentExistsError,
-    StackExistsError,
-)
-from zenml.integrations.neptune.experiment_trackers.run_state import get_neptune_run
-from zenml.integrations.neptune.flavors import NeptuneExperimentTrackerSettings
-from zenml.pipelines import pipeline
-from zenml.steps import step
 
 import neptune.new as neptune
 from tests.e2e.base import BaseE2ETest
 
+zenml = pytest.importorskip("zenml")
+pipelines = pytest.importorskip("zenml.pipelines")
+
 NEPTUNE_EXPERIMENT_TRACKER_NAME = "neptune_tracker"
 NEPTUNE_STACK_NAME = "neptune_stack"
 
 
-@fixture(scope="session")
+@pytest.fixture(scope="session")
 def registered_stack(zenml_client, experiment_tracker_comp, stack_with_neptune):
     try:
         zenml_client.initialize()
-    except InitializationException:
+    except zenml.exceptions.InitializationException:
         pass
 
     zenml_client.activate_stack(NEPTUNE_STACK_NAME)
 
 
-@fixture(scope="session")
-def zenml_client() -> Client:
+@pytest.fixture(scope="session")
+def zenml_client():
+    from zenml.client import Client
+
     return Client()
 
 
-@fixture(scope="session")
+@pytest.fixture(scope="session")
 def experiment_tracker_comp(zenml_client, environment):
     try:
         return zenml_client.create_stack_component(
             name=NEPTUNE_EXPERIMENT_TRACKER_NAME,
-            component_type=StackComponentType.EXPERIMENT_TRACKER,
+            component_type=zenml.enums.StackComponentType.EXPERIMENT_TRACKER,
             flavor="neptune",
             configuration={"api_token": environment.user_token, "project": environment.project},
         )
-    except StackComponentExistsError:
+    except zenml.exceptions.StackComponentExistsError:
         return zenml_client.get_stack_component(
-            component_type=StackComponentType.EXPERIMENT_TRACKER, name_id_or_prefix=NEPTUNE_EXPERIMENT_TRACKER_NAME
+            component_type=zenml.enums.StackComponentType.EXPERIMENT_TRACKER,
+            name_id_or_prefix=NEPTUNE_EXPERIMENT_TRACKER_NAME,
         )
 
 
-@fixture(scope="session")
+@pytest.fixture(scope="session")
 def stack_with_neptune(zenml_client, experiment_tracker_comp):
     a_id = zenml_client.active_stack.artifact_store.id
     o_id = zenml_client.active_stack.orchestrator.id
@@ -66,58 +58,57 @@ def stack_with_neptune(zenml_client, experiment_tracker_comp):
         return zenml_client.create_stack(
             name=NEPTUNE_STACK_NAME,
             components={
-                StackComponentType.ARTIFACT_STORE: a_id,
-                StackComponentType.ORCHESTRATOR: o_id,
-                StackComponentType.EXPERIMENT_TRACKER: e_id,
+                zenml.enums.StackComponentType.ARTIFACT_STORE: a_id,
+                zenml.enums.StackComponentType.ORCHESTRATOR: o_id,
+                zenml.enums.StackComponentType.EXPERIMENT_TRACKER: e_id,
             },
         )
-    except StackExistsError:
+    except zenml.exceptions.StackExistsError:
         return zenml_client.get_stack(name_id_or_prefix=NEPTUNE_STACK_NAME)
 
 
-settings = NeptuneExperimentTrackerSettings(tags={"sklearn", "digits"})
-
-
-@step(
-    experiment_tracker=NEPTUNE_EXPERIMENT_TRACKER_NAME,
-    settings={"experiment_tracker.neptune": settings},
-    enable_cache=False,
-)
-def example_step() -> None:
-    """A very minimalistic pipeline step.
-    Loads a sample dataset, trains a simple classifier and logs
-    a couple of metrics.
-    """
-    neptune_run = get_neptune_run()
-    digits = load_digits()
-    data = digits.images.reshape((len(digits.images), -1))
-
-    x_train, x_test, y_train, y_test = train_test_split(data, digits.target, test_size=0.3)
-    gamma = 0.001
-    neptune_run["params/gamma"] = gamma
-    model = SVC(gamma=gamma)
-    model.fit(x_train, y_train)
-    test_acc = model.score(x_test, y_test)
-    neptune_run["metrics/val_accuracy"] = test_acc
-
-
-@pipeline(
-    enable_cache=False,
-)
-def neptune_example_pipeline(ex_step):
-    """
-    Link all the steps artifacts together
-    """
-    ex_step()
-
-
-@mark.integrations
-@mark.zenml
+@pytest.mark.integrations
+@pytest.mark.zenml
 class TestZenML(BaseE2ETest):
     def _test_setup_creates_stack_with_neptune_experiment_tracker(self, zenml_client):
         assert zenml_client.active_stack.experiment_tracker.name == NEPTUNE_EXPERIMENT_TRACKER_NAME
 
     def _test_pipeline_runs_without_errors(self):
+        @zenml.steps.step(
+            experiment_tracker=NEPTUNE_EXPERIMENT_TRACKER_NAME,
+            settings={
+                "experiment_tracker.neptune": zenml.integrations.neptune.flavors.NeptuneExperimentTrackerSettings(
+                    tags={"sklearn", "digits"}
+                )
+            },
+            enable_cache=False,
+        )
+        def example_step() -> None:
+            """A very minimalistic pipeline step.
+            Loads a sample dataset, trains a simple classifier and logs
+            a couple of metrics.
+            """
+            neptune_run = zenml.integrations.neptune.experiment_trackers.run_state.get_neptune_run()
+            digits = load_digits()
+            data = digits.images.reshape((len(digits.images), -1))
+
+            x_train, x_test, y_train, y_test = train_test_split(data, digits.target, test_size=0.3)
+            gamma = 0.001
+            neptune_run["params/gamma"] = gamma
+            model = SVC(gamma=gamma)
+            model.fit(x_train, y_train)
+            test_acc = model.score(x_test, y_test)
+            neptune_run["metrics/val_accuracy"] = test_acc
+
+        @pipelines.pipeline(
+            enable_cache=False,
+        )
+        def neptune_example_pipeline(ex_step):
+            """
+            Link all the steps artifacts together
+            """
+            ex_step()
+
         run = neptune_example_pipeline(ex_step=example_step())
         run.run()