def test_create_delete_with_association(sagemaker_session):
obj = context.Context(sagemaker_session,
context_name="foo",
context_arn="foo")
sagemaker_session.sagemaker_client.list_associations.side_effect = [
{
"AssociationSummaries": [{
"SourceArn": obj.context_arn,
"SourceName": "X" + str(i),
"DestinationArn": "B" + str(i),
"DestinationName": "Y" + str(i),
"SourceType": "C" + str(i),
"DestinationType": "D" + str(i),
"AssociationType": "E" + str(i),
"CreationTime": None,
"CreatedBy": {},
} for i in range(1)],
},
{
"AssociationSummaries": [{
"SourceArn": "A" + str(i),
"SourceName": "X" + str(i),
"DestinationArn": obj.context_arn,
"DestinationName": "Y" + str(i),
"SourceType": "C" + str(i),
"DestinationType": "D" + str(i),
"AssociationType": "E" + str(i),
"CreationTime": None,
"CreatedBy": {},
} for i in range(1, 2)]
},
]
sagemaker_session.sagemaker_client.delete_association.return_value = {}
sagemaker_session.sagemaker_client.delete_context.return_value = {}
obj.delete(disassociate=True)
delete_with_association_expected_calls = [
unittest.mock.call(SourceArn=obj.context_arn, DestinationArn="B0"),
unittest.mock.call(SourceArn="A1", DestinationArn=obj.context_arn),
]
assert (delete_with_association_expected_calls ==
sagemaker_session.sagemaker_client.delete_association.mock_calls)
def test_save(sagemaker_session):
obj = context.Context(
sagemaker_session,
context_name="foo",
description="test-description",
properties={
"k1": "v1",
"k2": "v2"
},
properties_to_remove=["E"],
)
sagemaker_session.sagemaker_client.update_context.return_value = {}
obj.save()
sagemaker_session.sagemaker_client.update_context.assert_called_with(
ContextName="foo",
Description="test-description",
Properties={
"k1": "v1",
"k2": "v2"
},
PropertiesToRemove=["E"],
)
def test_delete(sagemaker_session):
obj = context.Context(sagemaker_session, context_name="foo")
sagemaker_session.sagemaker_client.delete_context.return_value = {}
obj.delete()
sagemaker_session.sagemaker_client.delete_context.assert_called_with(
ContextName="foo")