def test_update_model_version(self, mock_put):
updated_model_version = self.model_version_response.copy()
updated_model_version["data"]["items"][0]["description"] = "blabla"
mock_put.return_value = MockReturnValue(
json_data=updated_model_version)
res = models.update_model_version(model_id=1,
version_id=1,
description="blabla")
assert isinstance(res, ModelVersionResponse)
assert res.description == "blabla"
def check_gzip_enabled_and_return_mock(arg_url, data=None, headers=None, params=None, cookies=None):
# URL is sent as is
assert arg_url == url
# gzip is added as Content-Encoding header
assert headers["Content-Encoding"] == "gzip"
# data is gzipped. Decompress and check if items size matches
decompressed_assets = json.loads(gzip.decompress(data))
assert len(decompressed_assets["data"]["items"]) == len(RESPONSE)
# Return the mock response
return MockReturnValue(json_data=RESPONSE)
def test_create_version(self, post_mock):
model_version_response = self.model_version_response.copy()
model_version_response["data"]["items"][0]["trainingDetails"] = None
post_mock.return_value = MockReturnValue(
json_data=model_version_response)
res = models.create_model_version(model_id=1,
name="mymodel",
source_package_id=1)
assert isinstance(res, ModelVersionResponse)
assert 1 == res.id
def test_predict_instance_is_data_spec(self, mock_put, mock_data_spec):
mock_put.return_value = MockReturnValue(
json_data={"data": {
"predictions": [1, 2, 3]
}})
models.online_predict(model_id=1,
version_id=1,
instances=[mock_data_spec, mock_data_spec])
data_sent_to_api = json.loads(
gzip.decompress(mock_put.call_args[1]["data"]).decode())
for instance in data_sent_to_api["instances"]:
assert {"spec": "spec"} == instance
def test_put_request_failed(self, mock_request, api_client, url):
mock_request.return_value = MockReturnValue(
status=400, json_data={"error": "Client error"})
with pytest.raises(APIError) as e:
api_client._put(url)
assert re.match("Client error", str(e.value))
mock_request.return_value = MockReturnValue(status=500,
content="Server error")
with pytest.raises(APIError) as e:
api_client._put(url)
assert re.match("Server error", str(e.value))
mock_request.return_value = MockReturnValue(
status=500, json_data={"error": "Server error"})
with pytest.raises(APIError) as e:
api_client._put(url)
assert re.match("Server error", str(e.value))
def test_create_and_deploy_model_version(self, post_mock, put_mock):
model_version_response = self.model_version_response.copy()
model_version_response["data"]["items"][0]["trainingDetails"] = None
post_mock.side_effect = [
MockReturnValue(json_data=model_version_response),
MockReturnValue(
json_data={"data": {
"uploadUrl": "https://upload.here"
}}),
MockReturnValue(
json_data={"data": {
"uploadUrl": "https://upload.here"
}}),
MockReturnValue(json_data=model_version_response),
]
put_mock.return_value = MockReturnValue()
artifacts_directory = os.path.join(
os.path.dirname(__file__), "source_package_for_tests/artifacts")
model_version = models.deploy_model_version(
model_id=1,
name="mymodel",
source_package_id=1,
artifacts_directory=artifacts_directory)
assert model_version.id == 1
assert {
"description": "",
"metadata": {},
"name": "mymodel",
"sourcePackageId": 1,
} == get_call_args_data_from_mock(post_mock, 0, decompress_gzip=True)
post_artifacts_call_args = [
get_call_args_data_from_mock(post_mock, 1, decompress_gzip=True),
get_call_args_data_from_mock(post_mock, 2, decompress_gzip=True),
]
assert {"name": "artifact1.txt"} in post_artifacts_call_args
assert {"name": "sub_dir/artifact2.txt"} in post_artifacts_call_args
assert {} == get_call_args_data_from_mock(post_mock,
3,
decompress_gzip=True)
def test_get_request_with_autopaging(self, mock_request, api_client, url):
mock_request.side_effect = [
MockReturnValue(
json_data={"data": {
"items": [1, 2, 3],
"nextCursor": "next"
}}),
MockReturnValue(
json_data={"data": {
"items": [4, 5, 6],
"nextCursor": "next"
}}),
MockReturnValue(
json_data={"data": {
"items": [7, 8, 9],
"nextCursor": None
}}),
]
res = api_client._get(url, params={}, autopaging=True)
assert mock_request.call_count == 3
assert {"data": {"items": [1, 2, 3, 4, 5, 6, 7, 8, 9]}} == res.json()
def check_gzip_disabled_and_return_mock(arg_url,
data=None,
headers=None,
params=None,
cookies=None,
timeout=None):
# URL is sent as is
assert arg_url == api_client._base_url + url
# gzip is not added as Content-Encoding header
assert "Content-Encoding" not in headers
# data is not gzipped.
assert len(json.loads(data)["data"]["items"]) == len(RESPONSE)
# Return the mock response
return MockReturnValue(json_data=RESPONSE)
def check_args_to_put_and_return_mock(arg_url,
data=None,
headers=None,
params=None,
cookies=None,
timeout=None):
# URL is sent as is
assert arg_url == api_client._base_url + url
# data is json encoded
assert len(json.loads(data)["data"]["items"]) == len(RESPONSE)
# cookies should be the same
assert cookies == {"a-cookie": "a-cookie-val"}
# Return the mock response
return MockReturnValue(json_data=RESPONSE)
def test_train_and_deploy_version_data_spec_arg(self, post_mock,
mock_data_spec):
post_mock.return_value = MockReturnValue(
json_data=self.model_version_response)
models.train_and_deploy_model_version(
model_id=1,
name="mymodel",
source_package_id=1,
args={"data_spec": mock_data_spec})
data_sent_to_api = json.loads(
gzip.decompress(post_mock.call_args[1]["data"]).decode())
assert {
"spec": "spec"
} == data_sent_to_api["trainingDetails"]["args"]["data_spec"]
def test_list_versions(self, get_mock):
get_mock.return_value = MockReturnValue(
json_data=self.model_version_response)
res = models.list_model_versions(model_id=1)
assert len(res) > 0
assert isinstance(res, ModelVersionCollectionResponse)
assert isinstance(res[:1], ModelVersionCollectionResponse)
assert isinstance(res[0], ModelVersionResponse)
for model in res:
assert isinstance(model, ModelVersionResponse)
assert res[0].to_json(
) == self.model_version_response["data"]["items"][0]
assert res[0].id == self.model_version_response["data"]["items"][0][
"id"]
def check_args_to_post_and_return_mock(arg_url,
data=None,
headers=None,
params=None,
cookies=None,
timeout=None):
# URL is sent as is
assert arg_url == api_client._base_url + url
# cookies should be the same
assert cookies == {"a-cookie": "a-cookie-val"}
# Return the mock response
return MockReturnValue(json_data=RESPONSE)
def test_list_artifacts(self, get_mock):
get_mock.return_value = MockReturnValue(
json_data={"data": {
"items": [{
"name": "a1",
"size": 1
}]
}})
res = models.list_artifacts(model_id=1, version_id=1)
assert len(res) > 0
assert isinstance(res, ModelArtifactCollectionResponse)
assert isinstance(res[:1], ModelArtifactCollectionResponse)
assert isinstance(res[0], ModelArtifactResponse)
assert res[0].name == "a1"
assert res[0].size == 1
def test_get_model_version_log(self, mock_get):
mock_get.return_value = MockReturnValue(json_data={
"data": {
"predict": ["l1", "l2", "l3"],
"train": ["l1", "l2", "l3"]
}
})
res = models.get_logs(model_id=1, version_id=1, log_type="both")
assert isinstance(res, ModelLogResponse)
assert res.prediction_logs == ["l1", "l2", "l3"]
assert res.training_logs == ["l1", "l2", "l3"]
assert (
res.__str__() ==
'{\n "predict": [\n "l1",\n "l2",\n "l3"\n ],\n '
'"train": [\n "l1",\n "l2",\n "l3"\n ]\n}')
def test_create_schedule_with_data_spec_objects(self, mock_post,
mock_data_spec):
mock_post.return_value = MockReturnValue(
json_data=self.schedule_response)
res = schedules.create_schedule(model_id=123,
name="myschedule",
schedule_data_spec=mock_data_spec,
args={"k": "v"},
metadata={"k": "v"})
assert isinstance(res, ScheduleResponse)
assert res.id == 123
data_sent_to_api = json.loads(
gzip.decompress(mock_post.call_args[1]["data"]).decode())
actual_data_spec = data_sent_to_api["dataSpec"]
assert {"spec": "spec"} == actual_data_spec
def test_delete_request_ok(self, mock_request, api_client, url):
mock_request.return_value = MockReturnValue(json_data=RESPONSE)
response = api_client._delete(url)
assert response.status_code == 200
assert len(response.json()["data"]["items"]) == len(RESPONSE)
def test_delete_version(self, delete_mock):
delete_mock.return_value = MockReturnValue()
res = models.delete_model_version(model_id=1, version_id=1)
assert res is None
def test_deprecate_model_version(self, mock_put):
mock_put.return_value = MockReturnValue(
json_data=self.model_version_response)
res = models.deprecate_model_version(model_id=1, version_id=1)
assert isinstance(res, ModelVersionResponse)
assert res.is_deprecated is True
def test_get_request(self, mock_request):
mock_request.return_value = MockReturnValue(json_data=RESPONSE)
response = utils.get_request(url)
assert response.status_code == 200
assert len(response.json()['data']['items']) == len(RESPONSE)
def test_delete_schedule(self, mock_delete):
mock_delete.return_value = MockReturnValue()
res = schedules.delete_schedule(id=1)
assert res is None
def test_get_schedule(self, mock_get):
mock_get.return_value = MockReturnValue(
json_data=self.schedule_response)
res = schedules.get_schedule(id=1)
assert isinstance(res, ScheduleResponse)
assert self.schedule_response["data"]["items"][0]["name"] == res.name