change: Clean up Pipeline unit tests (#3578)

Author: qidewenwhen

tests/unit/sagemaker/workflow/conftest.py

@@ -16,17 +16,17 @@
 
 import pytest
 
-from sagemaker import Session
 from sagemaker.workflow.pipeline_context import PipelineSession
 
 REGION = "us-west-2"
 BUCKET = "my-bucket"
 ROLE = "DummyRole"
 IMAGE_URI = "fakeimage"
+INSTANCE_TYPE = "ml.m4.xlarge"
 
 
 @pytest.fixture(scope="module")
-def client():
+def mock_client():
     """Mock client.
 
     Considerations when appropriate:
@@ -38,11 +38,12 @@ def client():
     client_mock._client_config.user_agent = (
         "Boto3/1.14.24 Python/3.8.5 Linux/5.4.0-42-generic Botocore/1.17.24 Resource"
     )
+    client_mock.describe_model.return_value = {"PrimaryContainer": {}, "Containers": {}}
     return client_mock
 
 
 @pytest.fixture(scope="module")
-def boto_session(client):
+def mock_boto_session(client):
     role_mock = Mock()
     type(role_mock).arn = PropertyMock(return_value=ROLE)
 
@@ -57,19 +58,9 @@ def boto_session(client):
 
 
 @pytest.fixture(scope="module")
-def pipeline_session(boto_session, client):
+def pipeline_session(mock_boto_session, mock_client):
     return PipelineSession(
-        boto_session=boto_session,
-        sagemaker_client=client,
-        default_bucket=BUCKET,
-    )
-
-
-@pytest.fixture(scope="module")
-def sagemaker_session(boto_session, client):
-    return Session(
-        boto_session=boto_session,
-        sagemaker_client=client,
-        sagemaker_runtime_client=client,
+        boto_session=mock_boto_session,
+        sagemaker_client=mock_client,
         default_bucket=BUCKET,
     )

tests/unit/sagemaker/workflow/test_automl_step.py

@@ -15,57 +15,14 @@
 import json
 
 import pytest
-from mock import Mock, PropertyMock
 from sagemaker.automl.automl import AutoML, AutoMLInput
 from sagemaker.exceptions import AutoMLStepInvalidModeError
 from sagemaker.workflow import ParameterString
 
 from sagemaker.workflow.automl_step import AutoMLStep
 from sagemaker.workflow.model_step import ModelStep
 from sagemaker.workflow.pipeline import Pipeline
-from sagemaker.workflow.pipeline_context import PipelineSession
-
-REGION = "us-west-2"
-BUCKET = "my-bucket"
-ROLE = "DummyRole"
-
-
-@pytest.fixture
-def client():
-    """Mock client.
-    Considerations when appropriate:
-        * utilize botocore.stub.Stubber
-        * separate runtime client from client
-    """
-    client_mock = Mock()
-    client_mock._client_config.user_agent = (
-        "Boto3/1.14.24 Python/3.8.5 Linux/5.4.0-42-generic Botocore/1.17.24 Resource"
-    )
-    return client_mock
-
-
-@pytest.fixture
-def boto_session(client):
-    role_mock = Mock()
-    type(role_mock).arn = PropertyMock(return_value=ROLE)
-
-    resource_mock = Mock()
-    resource_mock.Role.return_value = role_mock
-
-    session_mock = Mock(region_name=REGION)
-    session_mock.resource.return_value = resource_mock
-    session_mock.client.return_value = client
-
-    return session_mock
-
-
-@pytest.fixture
-def pipeline_session(boto_session, client):
-    return PipelineSession(
-        boto_session=boto_session,
-        sagemaker_client=client,
-        default_bucket=BUCKET,
-    )
+from tests.unit.sagemaker.workflow.conftest import ROLE
 
 
 def test_single_automl_step(pipeline_session):

tests/unit/sagemaker/workflow/test_condition_step.py

@@ -13,8 +13,6 @@
 from __future__ import absolute_import
 import json
 
-import pytest
-from mock import Mock, MagicMock
 from sagemaker.workflow.conditions import (
     ConditionEquals,
     ConditionGreaterThan,
@@ -33,20 +31,6 @@
 from tests.unit.sagemaker.workflow.helpers import CustomStep, ordered
 
 
-@pytest.fixture()
-def sagemaker_session():
-    boto_mock = Mock(name="boto_session", region_name="us-west-2")
-    session_mock = MagicMock(
-        name="sagemaker_session",
-        boto_session=boto_mock,
-        boto_region_name="us-west-2",
-        config=None,
-        local_mode=False,
-        account_id=Mock(),
-    )
-    return session_mock
-
-
 def test_condition_step():
     param = ParameterInteger(name="MyInt")
     cond = ConditionEquals(left=param, right=1)

tests/unit/sagemaker/workflow/test_emr_step.py

@@ -14,30 +14,13 @@
 
 import json
 
-import pytest
-
-from mock import Mock
-
 from sagemaker.workflow.emr_step import EMRStep, EMRStepConfig
 from sagemaker.workflow.steps import CacheConfig
 from sagemaker.workflow.pipeline import Pipeline, PipelineGraph
 from sagemaker.workflow.parameters import ParameterString
 from tests.unit.sagemaker.workflow.helpers import CustomStep, ordered
 
 
-@pytest.fixture()
-def sagemaker_session():
-    boto_mock = Mock(name="boto_session", region_name="us-west-2")
-    session_mock = Mock(
-        name="sagemaker_session",
-        boto_session=boto_mock,
-        boto_region_name="us-west-2",
-        config=None,
-        local_mode=False,
-    )
-    return session_mock
-
-
 def test_emr_step_with_one_step_config(sagemaker_session):
     emr_step_config = EMRStepConfig(
         jar="s3:/script-runner/script-runner.jar",

tests/unit/sagemaker/workflow/test_lambda_step.py

@@ -26,20 +26,6 @@
 from tests.unit.sagemaker.workflow.helpers import CustomStep, ordered
 
 
-@pytest.fixture()
-def sagemaker_session():
-    boto_mock = Mock(name="boto_session", region_name="us-west-2")
-    session_mock = MagicMock(
-        name="sagemaker_session",
-        boto_session=boto_mock,
-        boto_region_name="us-west-2",
-        config=None,
-        local_mode=False,
-        account_id=Mock(),
-    )
-    return session_mock
-
-
 @pytest.fixture()
 def sagemaker_session_cn():
     boto_mock = Mock(name="boto_session", region_name="cn-north-1")