Init tests and stubs

metadata-ingestion/tests/unit/test_sagemaker_source.py

@@ -0,0 +1,69 @@
+import json
+
+from botocore.stub import Stubber
+from freezegun import freeze_time
+
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.source.glue import (
+    SagemakerSource,
+    SagemakerSourceConfig,
+    get_column_type,
+)
+from datahub.metadata.com.linkedin.pegasus2avro.schema import (
+    ArrayTypeClass,
+    MapTypeClass,
+    SchemaFieldDataType,
+    StringTypeClass,
+)
+from tests.test_helpers import mce_helpers
+from tests.unit.test_sagemaker_source_stubs import (
+    describe_feature_group_response_1,
+    describe_feature_group_response_2,
+    list_feature_groups_response,
+)
+
+FROZEN_TIME = "2020-04-14 07:00:00"
+
+
+def sagemaker_source() -> SagemakerSource:
+    return SagemakerSource(
+        ctx=PipelineContext(run_id="glue-source-test"),
+        config=SagemakerSourceConfig(aws_region="us-west-2"),
+    )
+
+
+@freeze_time(FROZEN_TIME)
+def test_sagemaker_ingest(tmp_path, pytestconfig):
+
+    sagemaker_source_instance = sagemaker_source()
+
+    with Stubber(sagemaker_source_instance.sagemaker_client) as sagemaker_stubber:
+
+        sagemaker_stubber.add_response(
+            "list_feature_groups", list_feature_groups_response, {}
+        )
+        sagemaker_stubber.add_response(
+            "describe_feature_group",
+            describe_feature_group_response_1,
+            {"FeatureGroupName": "test"},
+        )
+        sagemaker_stubber.add_response(
+            "describe_feature_group",
+            describe_feature_group_response_2,
+            {"FeatureGroupName": "test-1"},
+        )
+
+        mce_objects = [
+            wu.mce.to_obj() for wu in sagemaker_source_instance.get_workunits()
+        ]
+
+        with open(str(tmp_path / "sagemaker_mces.json"), "w") as f:
+            json.dump(mce_objects, f, indent=2)
+
+    output = mce_helpers.load_json_file(str(tmp_path / "sagemaker_mces.json"))
+
+    test_resources_dir = pytestconfig.rootpath / "tests/unit/sagemaker"
+    golden = mce_helpers.load_json_file(
+        str(test_resources_dir / "sagemaker_mces_golden.json")
+    )
+    mce_helpers.assert_mces_equal(output, golden)

metadata-ingestion/tests/unit/test_sagemaker_source_stubs.py

@@ -0,0 +1,57 @@
+import datetime
+
+list_feature_groups_response = {
+    "FeatureGroupSummaries": [
+        {
+            "FeatureGroupName": "test-1",
+            "FeatureGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test-1",
+            "CreationTime": datetime.datetime(2021, 6, 23, 13, 58, 10, 264000),
+            "FeatureGroupStatus": "Created",
+        },
+        {
+            "FeatureGroupName": "test",
+            "FeatureGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test",
+            "CreationTime": datetime.datetime(2021, 6, 14, 11, 3, 0, 803000),
+            "FeatureGroupStatus": "Created",
+        },
+    ],
+}
+
+describe_feature_group_response_1 = {
+    "FeatureGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test",
+    "FeatureGroupName": "test",
+    "RecordIdentifierFeatureName": "feature_1",
+    "EventTimeFeatureName": "feature_3",
+    "FeatureDefinitions": [
+        {"FeatureName": "feature_1", "FeatureType": "String"},
+        {"FeatureName": "feature_2", "FeatureType": "Integral"},
+        {"FeatureName": "feature_3", "FeatureType": "Fractional"},
+    ],
+    "CreationTime": datetime.datetime(
+        2021,
+        6,
+        14,
+        11,
+        3,
+        0,
+        803000,
+    ),
+    "OnlineStoreConfig": {"EnableOnlineStore": True},
+    "FeatureGroupStatus": "Created",
+}
+describe_feature_group_response_2 = {
+    "FeatureGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test-1",
+    "FeatureGroupName": "test-1",
+    "RecordIdentifierFeatureName": "id",
+    "EventTimeFeatureName": "time",
+    "FeatureDefinitions": [
+        {"FeatureName": "name", "FeatureType": "String"},
+        {"FeatureName": "id", "FeatureType": "Integral"},
+        {"FeatureName": "height", "FeatureType": "Fractional"},
+        {"FeatureName": "time", "FeatureType": "String"},
+    ],
+    "CreationTime": datetime.datetime(2021, 6, 23, 13, 58, 10, 264000),
+    "OnlineStoreConfig": {"EnableOnlineStore": True},
+    "FeatureGroupStatus": "Created",
+    "Description": "First test feature group",
+}