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

Co-authored-by: Ajay Karpur <[email protected]>

Author: danabens

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

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):

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)
     ]

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

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

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