Skip to content

fix: map user context is list associations response #2238

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 2, 2021
Merged

fix: map user context is list associations response #2238

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8d12105
fix: map user context is list associations response
danabens Mar 24, 2021
59807ec
Merge branch 'master' into user-context-fix
ajaykarpur Apr 1, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 33 additions & 1 deletion src/sagemaker/lineage/_api_types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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
Expand All @@ -204,4 +237,3 @@ class AssociationSummary(_base_types.ApiObject):
destination_type = None
association_type = None
creation_time = None
created_by = None
3 changes: 3 additions & 0 deletions tests/integ/sagemaker/lineage/test_association.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
18 changes: 15 additions & 3 deletions tests/unit/sagemaker/lineage/test_association.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
],
Expand All @@ -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)
]
Expand All @@ -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)
]
Expand Down
6 changes: 3 additions & 3 deletions tests/unit/sagemaker/lineage/test_dataset_artifact.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ def test_trained_models(sagemaker_session):
"DestinationType": "ModelDeployment",
"AssociationType": "E1",
"CreationTime": None,
"CreatedBy": {},
"CreatedBy": None,
}
],
},
Expand All @@ -57,7 +57,7 @@ def test_trained_models(sagemaker_session):
"DestinationType": "Context",
"AssociationType": "E2",
"CreationTime": None,
"CreatedBy": {},
"CreatedBy": None,
}
]
},
Expand All @@ -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
6 changes: 3 additions & 3 deletions tests/unit/sagemaker/lineage/test_endpoint_context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def test_models(sagemaker_session):
"DestinationType": "ModelDeployment",
"AssociationType": "E1",
"CreationTime": None,
"CreatedBy": {},
"CreatedBy": None,
}
],
},
Expand All @@ -47,7 +47,7 @@ def test_models(sagemaker_session):
"DestinationType": "Model",
"AssociationType": "E2",
"CreationTime": None,
"CreatedBy": {},
"CreatedBy": None,
}
]
},
Expand All @@ -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
6 changes: 3 additions & 3 deletions tests/unit/sagemaker/lineage/test_model_artifact.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ def test_trained_models(sagemaker_session):
"DestinationType": "Action",
"AssociationType": "E1",
"CreationTime": None,
"CreatedBy": {},
"CreatedBy": None,
}
],
},
Expand All @@ -49,7 +49,7 @@ def test_trained_models(sagemaker_session):
"DestinationType": "Context",
"AssociationType": "E2",
"CreationTime": None,
"CreatedBy": {},
"CreatedBy": None,
}
]
},
Expand All @@ -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