fix: map user context is list associations response (#2238)

danabens · ajaykarpur · web-flow · commit b66cb98b9223 · 2021-04-02T10:57:15.000-07:00
Co-authored-by: Ajay Karpur &lt;akarpur@amazon.com&gt;
diff --git a/src/sagemaker/lineage/_api_types.py b/src/sagemaker/lineage/_api_types.py
@@ -181,6 +181,36 @@ class ContextSummary(_base_types.ApiObject):
     last_modified_time = None
 
 
+class UserContext(_base_types.ApiObject):
+    """Summary model of a user context.
+
+    Attributes:
+        user_profile_arn (str): User profile ARN.
+        user_profile_name (str): User profile name.
+        domain_id (str): DomainId.
+    """
+
+    user_profile_arn = None
+    user_profile_name = None
+    domain_id = None
+
+    def __init__(self, user_profile_arn=None, user_profile_name=None, domain_id=None, **kwargs):
+        """Initialize UserContext.
+
+        Args:
+            user_profile_arn (str): User profile ARN.
+            user_profile_name (str): User profile name.
+            domain_id (str): DomainId.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super(UserContext, self).__init__(
+            user_profile_arn=user_profile_arn,
+            user_profile_name=user_profile_name,
+            domain_id=domain_id,
+            **kwargs
+        )
+
+
 class AssociationSummary(_base_types.ApiObject):
     """Summary model of an association.
 
@@ -196,6 +226,9 @@ class AssociationSummary(_base_types.ApiObject):
         created_by (obj): Context on creator.
     """
 
+    _custom_boto_types = {
+        "created_by": (UserContext, False),
+    }
     source_arn = None
     source_name = None
     destination_arn = None
@@ -204,4 +237,3 @@ class AssociationSummary(_base_types.ApiObject):
     destination_type = None
     association_type = None
     creation_time = None
-    created_by = None
diff --git a/tests/integ/sagemaker/lineage/test_association.py b/tests/integ/sagemaker/lineage/test_association.py
@@ -56,6 +56,9 @@ def test_list(association_objs, sagemaker_session):
     # sanity check
     assert association_keys_listed
 
+    for listed_asscn in listed:
+        assert listed_asscn.created_by is None
+
 
 @pytest.mark.timeout(30)
 def test_set_tag(association_obj, sagemaker_session):
diff --git a/tests/unit/sagemaker/lineage/test_association.py b/tests/unit/sagemaker/lineage/test_association.py
@@ -92,7 +92,11 @@ def test_list(sagemaker_session):
                     "DestinationType": "D" + str(i),
                     "AssociationType": "E" + str(i),
                     "CreationTime": creation_time + datetime.timedelta(hours=i),
-                    "CreatedBy": {},
+                    "CreatedBy": {
+                        "UserProfileArn": "profileArn",
+                        "UserProfileName": "profileName",
+                        "DomainId": "domainId",
+                    },
                 }
                 for i in range(10)
             ],
@@ -109,7 +113,11 @@ def test_list(sagemaker_session):
                     "DestinationType": "D" + str(i),
                     "AssociationType": "E" + str(i),
                     "CreationTime": creation_time + datetime.timedelta(hours=i),
-                    "CreatedBy": {},
+                    "CreatedBy": {
+                        "UserProfileArn": "profileArn",
+                        "UserProfileName": "profileName",
+                        "DomainId": "domainId",
+                    },
                 }
                 for i in range(10, 20)
             ]
@@ -126,7 +134,11 @@ def test_list(sagemaker_session):
             destination_type="D" + str(i),
             association_type="E" + str(i),
             creation_time=creation_time + datetime.timedelta(hours=i),
-            created_by={},
+            created_by=_api_types.UserContext(
+                user_profile_arn="profileArn",
+                user_profile_name="profileName",
+                domain_id="domainId",
+            ),
         )
         for i in range(20)
     ]
diff --git a/tests/unit/sagemaker/lineage/test_dataset_artifact.py b/tests/unit/sagemaker/lineage/test_dataset_artifact.py
@@ -42,7 +42,7 @@ def test_trained_models(sagemaker_session):
                     "DestinationType": "ModelDeployment",
                     "AssociationType": "E1",
                     "CreationTime": None,
-                    "CreatedBy": {},
+                    "CreatedBy": None,
                 }
             ],
         },
@@ -57,7 +57,7 @@ def test_trained_models(sagemaker_session):
                     "DestinationType": "Context",
                     "AssociationType": "E2",
                     "CreationTime": None,
-                    "CreatedBy": {},
+                    "CreatedBy": None,
                 }
             ]
         },
@@ -79,7 +79,7 @@ def test_trained_models(sagemaker_session):
             destination_type="Context",
             association_type="E2",
             creation_time=None,
-            created_by={},
+            created_by=None,
         )
     ]
     assert expected_model_list == model_list
diff --git a/tests/unit/sagemaker/lineage/test_endpoint_context.py b/tests/unit/sagemaker/lineage/test_endpoint_context.py
@@ -32,7 +32,7 @@ def test_models(sagemaker_session):
                     "DestinationType": "ModelDeployment",
                     "AssociationType": "E1",
                     "CreationTime": None,
-                    "CreatedBy": {},
+                    "CreatedBy": None,
                 }
             ],
         },
@@ -47,7 +47,7 @@ def test_models(sagemaker_session):
                     "DestinationType": "Model",
                     "AssociationType": "E2",
                     "CreationTime": None,
-                    "CreatedBy": {},
+                    "CreatedBy": None,
                 }
             ]
         },
@@ -71,7 +71,7 @@ def test_models(sagemaker_session):
             destination_type="Model",
             association_type="E2",
             creation_time=None,
-            created_by={},
+            created_by=None,
         )
     ]
     assert expected_model_list == model_list
diff --git a/tests/unit/sagemaker/lineage/test_model_artifact.py b/tests/unit/sagemaker/lineage/test_model_artifact.py
@@ -34,7 +34,7 @@ def test_trained_models(sagemaker_session):
                     "DestinationType": "Action",
                     "AssociationType": "E1",
                     "CreationTime": None,
-                    "CreatedBy": {},
+                    "CreatedBy": None,
                 }
             ],
         },
@@ -49,7 +49,7 @@ def test_trained_models(sagemaker_session):
                     "DestinationType": "Context",
                     "AssociationType": "E2",
                     "CreationTime": None,
-                    "CreatedBy": {},
+                    "CreatedBy": None,
                 }
             ]
         },
@@ -71,7 +71,7 @@ def test_trained_models(sagemaker_session):
             destination_type="Context",
             association_type="E2",
             creation_time=None,
-            created_by={},
+            created_by=None,
         )
     ]
     assert expected_model_list == endpoint_context_list

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ def test_trained_models(sagemaker_session):`
`42`	`42`	`"DestinationType": "ModelDeployment",`
`43`	`43`	`"AssociationType": "E1",`
`44`	`44`	`"CreationTime": None,`
`45`		`- "CreatedBy": {},`
	`45`	`+ "CreatedBy": None,`
`46`	`46`	`}`
`47`	`47`	`],`
`48`	`48`	`},`
`@@ -57,7 +57,7 @@ def test_trained_models(sagemaker_session):`
`57`	`57`	`"DestinationType": "Context",`
`58`	`58`	`"AssociationType": "E2",`
`59`	`59`	`"CreationTime": None,`
`60`		`- "CreatedBy": {},`
	`60`	`+ "CreatedBy": None,`
`61`	`61`	`}`
`62`	`62`	`]`
`63`	`63`	`},`
`@@ -79,7 +79,7 @@ def test_trained_models(sagemaker_session):`
`79`	`79`	`destination_type="Context",`
`80`	`80`	`association_type="E2",`
`81`	`81`	`creation_time=None,`
`82`		`- created_by={},`
	`82`	`+ created_by=None,`
`83`	`83`	`)`
`84`	`84`	`]`
`85`	`85`	`assert expected_model_list == model_list`
Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ def test_models(sagemaker_session):`
`32`	`32`	`"DestinationType": "ModelDeployment",`
`33`	`33`	`"AssociationType": "E1",`
`34`	`34`	`"CreationTime": None,`
`35`		`- "CreatedBy": {},`
	`35`	`+ "CreatedBy": None,`
`36`	`36`	`}`
`37`	`37`	`],`
`38`	`38`	`},`
`@@ -47,7 +47,7 @@ def test_models(sagemaker_session):`
`47`	`47`	`"DestinationType": "Model",`
`48`	`48`	`"AssociationType": "E2",`
`49`	`49`	`"CreationTime": None,`
`50`		`- "CreatedBy": {},`
	`50`	`+ "CreatedBy": None,`
`51`	`51`	`}`
`52`	`52`	`]`
`53`	`53`	`},`
`@@ -71,7 +71,7 @@ def test_models(sagemaker_session):`
`71`	`71`	`destination_type="Model",`
`72`	`72`	`association_type="E2",`
`73`	`73`	`creation_time=None,`
`74`		`- created_by={},`
	`74`	`+ created_by=None,`
`75`	`75`	`)`
`76`	`76`	`]`
`77`	`77`	`assert expected_model_list == model_list`
Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@ def test_trained_models(sagemaker_session):`
`34`	`34`	`"DestinationType": "Action",`
`35`	`35`	`"AssociationType": "E1",`
`36`	`36`	`"CreationTime": None,`
`37`		`- "CreatedBy": {},`
	`37`	`+ "CreatedBy": None,`
`38`	`38`	`}`
`39`	`39`	`],`
`40`	`40`	`},`
`@@ -49,7 +49,7 @@ def test_trained_models(sagemaker_session):`
`49`	`49`	`"DestinationType": "Context",`
`50`	`50`	`"AssociationType": "E2",`
`51`	`51`	`"CreationTime": None,`
`52`		`- "CreatedBy": {},`
	`52`	`+ "CreatedBy": None,`
`53`	`53`	`}`
`54`	`54`	`]`
`55`	`55`	`},`
`@@ -71,7 +71,7 @@ def test_trained_models(sagemaker_session):`
`71`	`71`	`destination_type="Context",`
`72`	`72`	`association_type="E2",`
`73`	`73`	`creation_time=None,`
`74`		`- created_by={},`
	`74`	`+ created_by=None,`
`75`	`75`	`)`
`76`	`76`	`]`
`77`	`77`	`assert expected_model_list == endpoint_context_list`