def test_create_delete_with_association(sagemaker_session):
obj = action.Action(
sagemaker_session,
action_name="foo",
action_arn="a1",
description="updated-description",
status="updated-status",
properties={"k1": "v1"},
properties_to_remove=["k2"],
)
sagemaker_session.sagemaker_client.list_associations.side_effect = [
{
"AssociationSummaries": [{
"SourceArn": obj.action_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.action_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_action.return_value = {}
obj.delete(disassociate=True)
delete_with_association_expected_calls = [
unittest.mock.call(SourceArn=obj.action_arn, DestinationArn="B0"),
unittest.mock.call(SourceArn="A1", DestinationArn=obj.action_arn),
]
assert (delete_with_association_expected_calls ==
sagemaker_session.sagemaker_client.delete_association.mock_calls)
def test_save(sagemaker_session):
obj = action.Action(
sagemaker_session,
action_name="foo",
description="updated-description",
status="updated-status",
properties={"k1": "v1"},
properties_to_remove=["k2"],
)
sagemaker_session.sagemaker_client.update_action.return_value = {}
obj.save()
sagemaker_session.sagemaker_client.update_action.assert_called_with(
ActionName="foo",
Description="updated-description",
Status="updated-status",
Properties={"k1": "v1"},
PropertiesToRemove=["k2"],
)
def test_delete(sagemaker_session):
obj = action.Action(sagemaker_session, action_name="foo", source_uri="bar")
sagemaker_session.sagemaker_client.delete_action.return_value = {}
obj.delete()
sagemaker_session.sagemaker_client.delete_action.assert_called_with(
ActionName="foo")
