1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
|
import boto3
import pytest
from moto import mock_aws
from moto.core import DEFAULT_ACCOUNT_ID as ACCOUNT_ID
TEST_REGION_NAME = "us-east-1"
TEST_EXPERIMENT_NAME = "MyExperimentName"
@pytest.fixture(name="sagemaker_client")
def fixture_sagemaker_client():
with mock_aws():
yield boto3.client("sagemaker", region_name=TEST_REGION_NAME)
def test_create_experiment(sagemaker_client):
resp = sagemaker_client.create_experiment(ExperimentName=TEST_EXPERIMENT_NAME)
assert resp["ResponseMetadata"]["HTTPStatusCode"] == 200
resp = sagemaker_client.list_experiments()
assert len(resp["ExperimentSummaries"]) == 1
assert resp["ExperimentSummaries"][0]["ExperimentName"] == TEST_EXPERIMENT_NAME
assert (
resp["ExperimentSummaries"][0]["ExperimentArn"]
== f"arn:aws:sagemaker:{TEST_REGION_NAME}:{ACCOUNT_ID}:experiment/{TEST_EXPERIMENT_NAME}"
)
def test_list_experiments(sagemaker_client):
experiment_names = [f"some-experiment-name-{i}" for i in range(10)]
for experiment_name in experiment_names:
resp = sagemaker_client.create_experiment(ExperimentName=experiment_name)
assert resp["ResponseMetadata"]["HTTPStatusCode"] == 200
resp = sagemaker_client.list_experiments(MaxResults=1)
assert len(resp["ExperimentSummaries"]) == 1
next_token = resp["NextToken"]
resp = sagemaker_client.list_experiments(MaxResults=2, NextToken=next_token)
assert len(resp["ExperimentSummaries"]) == 2
next_token = resp["NextToken"]
resp = sagemaker_client.list_experiments(NextToken=next_token)
assert len(resp["ExperimentSummaries"]) == 7
assert resp.get("NextToken") is None
def test_delete_experiment(sagemaker_client):
sagemaker_client.create_experiment(ExperimentName=TEST_EXPERIMENT_NAME)
resp = sagemaker_client.delete_experiment(ExperimentName=TEST_EXPERIMENT_NAME)
assert resp["ResponseMetadata"]["HTTPStatusCode"] == 200
resp = sagemaker_client.list_experiments()
assert len(resp["ExperimentSummaries"]) == 0
def test_add_tags_to_experiment(sagemaker_client):
sagemaker_client.create_experiment(ExperimentName=TEST_EXPERIMENT_NAME)
resp = sagemaker_client.describe_experiment(ExperimentName=TEST_EXPERIMENT_NAME)
arn = resp["ExperimentArn"]
tags = [{"Key": "name", "Value": "value"}]
sagemaker_client.add_tags(ResourceArn=arn, Tags=tags)
assert resp["ResponseMetadata"]["HTTPStatusCode"] == 200
resp = sagemaker_client.list_tags(ResourceArn=arn)
assert resp["Tags"] == tags
def test_delete_tags_to_experiment(sagemaker_client):
sagemaker_client.create_experiment(ExperimentName=TEST_EXPERIMENT_NAME)
resp = sagemaker_client.describe_experiment(ExperimentName=TEST_EXPERIMENT_NAME)
arn = resp["ExperimentArn"]
tags = [{"Key": "name", "Value": "value"}]
sagemaker_client.add_tags(ResourceArn=arn, Tags=tags)
sagemaker_client.delete_tags(ResourceArn=arn, TagKeys=[i["Key"] for i in tags])
assert resp["ResponseMetadata"]["HTTPStatusCode"] == 200
resp = sagemaker_client.list_tags(ResourceArn=arn)
assert resp["Tags"] == []
|
