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"]