def test_upload_artifacts_from_directory(self, post_mock, put_mock):
artifacts_directory = os.path.join(
os.path.dirname(__file__), "source_package_for_tests/artifacts")
post_mock.side_effect = [
MockReturnValue(
json_data={"data": {
"uploadUrl": "https://upload.here"
}}),
MockReturnValue(
json_data={"data": {
"uploadUrl": "https://upload.here"
}}),
]
put_mock.return_value = MockReturnValue()
models.upload_artifacts_from_directory(model_id=1,
version_id=1,
directory=artifacts_directory)
post_artifacts_call_args = [
get_call_args_data_from_mock(post_mock, 0, decompress_gzip=True),
get_call_args_data_from_mock(post_mock, 1, decompress_gzip=True),
]
assert {"name": "artifact1.txt"} in post_artifacts_call_args
assert {"name": "sub_dir/artifact2.txt"} in post_artifacts_call_args
def test_delete_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._delete(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._delete(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._delete(url)
assert re.match("Server error", str(e.value))
mock_request.return_value = MockReturnValue(
status=400,
json_data={"error": {
"code": 400,
"message": "Client error"
}})
with pytest.raises(APIError) as e:
api_client._delete(url)
assert re.match("Client error | code: 400 | X-Request-ID:",
str(e.value))
assert e.value.code == 400
assert e.value.message == "Client error"
def test_upload_artifact_from_file(self, post_mock, put_mock,
artifact_file_path):
post_mock.return_value = MockReturnValue(
json_data={"data": {
"uploadUrl": "https://upload.here"
}})
put_mock.return_value = MockReturnValue()
models.upload_artifact_from_file(model_id=1,
version_id=1,
name="my_artifact.txt",
file_path=artifact_file_path)
def test_predict_on_model_version(self, mock_put):
mock_put.return_value = MockReturnValue(
json_data={"data": {
"predictions": [1, 2, 3]
}})
predictions = models.online_predict(model_id=1, version_id=1)
assert predictions == [1, 2, 3]
def test_train_and_deploy_version(self, post_mock):
post_mock.return_value = MockReturnValue(
json_data=self.model_version_response)
res = models.train_and_deploy_model_version(model_id=1,
name="mymodel",
source_package_id=1)
assert isinstance(res, ModelVersionResponse)
def tagmatching_result(mock_post):
response = {
"data": {
"items": [{
"matches": [
{
"platform": "a_platform",
"score": 0,
"tagId": "a_match"
},
{
"platform": "a_platform",
"score": 0,
"tagId": "a_match1"
},
{
"platform": "a_platform",
"score": 0,
"tagId": "a_match2"
},
],
"tagId":
"a_tag",
}]
}
}
mock_post.return_value = MockReturnValue(status=200, json_data=response)
return tag_matching(tag_ids=["a_tag"])
def tagmatching_result(mock_post):
from cognite.v04.tagmatching import tag_matching
response = {
'data': {
'items': [{
'matches': [{
'platform': 'a_platform',
'score': 0,
'tagId': 'a_match'
}, {
'platform': 'a_platform',
'score': 0,
'tagId': 'a_match1'
}, {
'platform': 'a_platform',
'score': 0,
'tagId': 'a_match2'
}],
'tagId':
'a_tag'
}]
}
}
mock_post.return_value = MockReturnValue(status=200, json_data=response)
return tag_matching(tag_ids=['a_tag'])
def test_get_version(self, get_mock):
get_mock.return_value = MockReturnValue(
json_data=self.model_version_response)
model_version = models.get_model_version(model_id=1, version_id=1)
assert isinstance(model_version, ModelVersionResponse)
assert model_version.id == self.model_version_response["data"][
"items"][0]["id"]
def test_put_request_exception(self, mock_request, api_client, url):
mock_request.return_value = MockReturnValue(status=500)
mock_request.side_effect = Exception("Custom error")
with pytest.raises(Exception) as e:
api_client._put(url)
assert re.match("Custom error", str(e.value))
def test_get_request_exception(self, mock_request):
mock_request.return_value = MockReturnValue(status=500)
mock_request.side_effect = Exception("Custom error")
with pytest.raises(utils.APIError) as e:
utils.get_request(url)
assert re.match("Custom error", str(e.value))
def test_deprecate_schedule(self, mock_put):
depr_schedule_response = self.schedule_response.copy()
depr_schedule_response["data"]["items"][0]["isDeprecated"] = True
mock_put.return_value = MockReturnValue(
json_data=depr_schedule_response)
res = schedules.deprecate_schedule(id=1)
assert res.is_deprecated is True
def test_download_artifact(self, mock_get):
mock_get.side_effect = [
MockReturnValue(
json_data={"data": {
"downloadUrl": "https://download.me"
}}),
MockReturnValue(content=b"content"),
]
models.download_artifact(model_id=1,
version_id=1,
name="a1",
directory=os.getcwd())
file_path = os.path.join(os.getcwd(), "a1")
assert os.path.isfile(file_path)
with open(file_path, "rb") as f:
assert b"content" == f.read()
os.remove(file_path)
def test_put_request_ok(self, mock_request, api_client, url):
mock_request.return_value = MockReturnValue(json_data=RESPONSE)
response = api_client._put(url, body=RESPONSE)
response_json = response.json()
assert response.status_code == 200
assert len(response_json["data"]["items"]) == len(RESPONSE)
def test_predict_on_model_version_prediction_error(self, mock_put):
mock_put.return_value = MockReturnValue(
json_data={"error": {
"message": "User error",
"code": 200
}})
with pytest.raises(PredictionError, match="User error"):
models.online_predict(model_id=1, version_id=1)
def test_post_request_ok(self, mock_request):
mock_request.return_value = MockReturnValue(json_data=RESPONSE)
response = utils.post_request(url, RESPONSE)
response_json = response.json()
assert response.status_code == 200
assert len(response_json['data']['items']) == len(RESPONSE)
def check_args_to_put_and_return_mock(arg_url, data=None, headers=None, params=None, cookies=None):
# URL is sent as is
assert arg_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 check_args_to_post_and_return_mock(arg_url, data=None, headers=None, params=None, cookies=None):
# URL is sent as is
assert arg_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_deploy_version(self, post_mock):
model_version_response = self.model_version_response.copy()
model_version_response["data"]["items"][0]["trainingDetails"] = None
model_version_response["data"]["items"][0]["status"] = "DEPLOYING"
post_mock.return_value = MockReturnValue(
json_data=model_version_response)
res = models.deploy_awaiting_model_version(model_id=1, version_id=1)
assert isinstance(res, ModelVersionResponse)
assert "DEPLOYING" == res.status
def test_get_request_failed(self, mock_request):
mock_request.return_value = MockReturnValue(status=400, json_data={"error": "Client error"})
with pytest.raises(utils.APIError) as e:
utils.get_request(url)
assert re.match("Client error[\n]X-Request_id", str(e.value))
mock_request.return_value = MockReturnValue(status=500, content="Server error")
with pytest.raises(utils.APIError) as e:
utils.get_request(url)
assert re.match("Server error[\n]X-Request_id", str(e.value))
mock_request.return_value = MockReturnValue(status=500, json_data={"error": "Server error"})
with pytest.raises(utils.APIError) as e:
utils.get_request(url)
assert re.match("Server error[\n]X-Request_id", str(e.value))
def check_gzip_disabled_and_return_mock(arg_url, data=None, headers=None, params=None, cookies=None):
# URL is sent as is
assert arg_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_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_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 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_list_schedules(self, mock_get):
mock_get.return_value = MockReturnValue(
json_data=self.schedule_response)
res = schedules.list_schedules(limit=1)
assert len(res) > 0
assert isinstance(res, ScheduleCollectionResponse)
assert isinstance(res[:1], ScheduleCollectionResponse)
assert isinstance(res[0], ScheduleResponse)
assert self.schedule_response["data"]["items"][0]["name"] == res[
0].name
def test_create_schedule(self, mock_post):
mock_post.return_value = MockReturnValue(
json_data=self.schedule_response)
res = schedules.create_schedule(model_id=123,
name="myschedule",
schedule_data_spec={"spec": "spec"},
args={"k": "v"},
metadata={"k": "v"})
assert isinstance(res, ScheduleResponse)
assert res.id == 123
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_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_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 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"]
