From f4a63372e9bb33408a00e8879b1f69aea385174c Mon Sep 17 00:00:00 2001 From: Erlend vollset Date: Mon, 15 Aug 2022 16:07:50 +0200 Subject: [PATCH] Some fixes and cleanups to functions/relationships/contextualization SDKs (#996) --- cognite/client/_api/functions.py | 102 ++++++++---------- cognite/client/_api/relationships.py | 4 +- cognite/client/data_classes/functions.py | 4 +- cognite/client/data_classes/relationships.py | 2 +- tests/tests_unit/test_api/test_functions.py | 89 +++++++-------- .../test_contextualization.py | 5 +- 6 files changed, 102 insertions(+), 104 deletions(-) diff --git a/cognite/client/_api/functions.py b/cognite/client/_api/functions.py index 1171734ec1..589a1bd707 100644 --- a/cognite/client/_api/functions.py +++ b/cognite/client/_api/functions.py @@ -434,26 +434,25 @@ def _zip_and_upload_folder(self, folder: str, name: str, external_id: Optional[s try: with TemporaryDirectory() as tmpdir: zip_path = os.path.join(tmpdir, "function.zip") - zf = ZipFile(zip_path, "w") - for root, dirs, files in os.walk("."): - zf.write(root) - - # Validate requirements.txt in root-dir only - if root == "." and REQUIREMENTS_FILE_NAME in files: - # Remove requirement from file list - path = files.pop(files.index(REQUIREMENTS_FILE_NAME)) - reqs = _extract_requirements_from_file(path) - # Validate and format requirements - req_path = _validate_requirements(reqs) - - # NOTE: the actual file is not written. - # A temporary formatted file is used instead - zf.write(req_path, arcname=REQUIREMENTS_FILE_NAME) - - for filename in files: - zf.write(os.path.join(root, filename)) - - zf.close() + with ZipFile(zip_path, "w") as zf: + for root, dirs, files in os.walk("."): + zf.write(root) + + # Validate requirements.txt in root-dir only + if root == "." and REQUIREMENTS_FILE_NAME in files: + # Remove requirement from file list + path = files.pop(files.index(REQUIREMENTS_FILE_NAME)) + reqs = _extract_requirements_from_file(path) + # Validate and format requirements + parsed_reqs = _validate_and_parse_requirements(reqs) + with NamedTemporaryFile() as nth: + _write_requirements_to_file(nth.name, parsed_reqs) + # NOTE: the actual file is not written. + # A temporary formatted file is used instead + zf.write(nth.name, arcname=REQUIREMENTS_FILE_NAME) + + for filename in files: + zf.write(os.path.join(root, filename)) overwrite = True if external_id else False file = cast( @@ -466,8 +465,6 @@ def _zip_and_upload_folder(self, folder: str, name: str, external_id: Optional[s file_id = cast(int, file.id) return file_id - except Exception as e: - raise e finally: os.chdir(current_dir) @@ -482,16 +479,16 @@ def _zip_and_upload_handle(self, function_handle: Callable, name: str, external_ f.write(source) # Read and validate requirements - req_path = _get_requirements_handle(fn=function_handle) + with NamedTemporaryFile() as named_temp_file: + requirements_written = _write_fn_docstring_requirements_to_file(function_handle, named_temp_file.name) - zip_path = os.path.join(tmpdir, "function.zip") - zf = ZipFile(zip_path, "w") - zf.write(handle_path, arcname=HANDLER_FILE_NAME) + zip_path = os.path.join(tmpdir, "function.zip") + with ZipFile(zip_path, "w") as zf: + zf.write(handle_path, arcname=HANDLER_FILE_NAME) - # Zip requirements.txt - if req_path: - zf.write(req_path, arcname=REQUIREMENTS_FILE_NAME) - zf.close() + # Zip requirements.txt + if requirements_written: + zf.write(named_temp_file.name, arcname=REQUIREMENTS_FILE_NAME) overwrite = True if external_id else False file = cast( @@ -602,19 +599,16 @@ def _use_token_exchange( raise CogniteAPIError("Failed to create session using token exchange flow.", 403) from e -def _using_client_credential_flow( - cognite_client: "CogniteClient", -) -> bool: +def _using_client_credential_flow(cognite_client: "CogniteClient") -> bool: """ Determine whether the Cognite client is configured for client-credential flow. """ client_config = cognite_client.config - return cast( - bool, - client_config.token_client_secret - and client_config.token_client_id - and client_config.token_url - and client_config.token_scopes, + return ( + client_config.token_client_secret is not None + and client_config.token_client_id is not None + and client_config.token_url is not None + and client_config.token_scopes is not None ) @@ -716,17 +710,15 @@ def _extract_requirements_from_doc_string(docstr: str) -> Optional[List[str]]: return None -def _validate_requirements(requirements: List[str]) -> str: +def _validate_and_parse_requirements(requirements: List[str]) -> List[str]: """Validates the requirement specifications Args: requirements (list[str]): list of requirement specifications - Raises: ValueError: if validation of requirements fails - Returns: - str: output path of the requirements file + List[str]: The parsed requirements """ parsed_reqs: List[str] = [] for req in requirements: @@ -736,34 +728,34 @@ def _validate_requirements(requirements: List[str]) -> str: raise ValueError(str(e)) parsed_reqs.append(str(parsed).strip()) + return parsed_reqs - tmp = NamedTemporaryFile() - # Write requirements to temporary file - with open(tmp.name, "w+") as f: - f.write("\n".join(parsed_reqs)) - f.close() +def _write_requirements_to_file(file_path: str, requirements: List[str]) -> None: + with open(file_path, "w+") as f: + f.write("\n".join(requirements)) - return tmp.name - -def _get_requirements_handle(fn: Callable) -> Optional[str]: - """Read requirements from a function docstring, and validate them +def _write_fn_docstring_requirements_to_file(fn: Callable, file_path: str) -> bool: + """Read requirements from a function docstring, validate them, and write contents to the provided file path Args: fn (Callable): the function to read requirements from + file_path (str): Path of file to write requirements to Returns: - str: output path of the requirements file, or None if no requirements are specified + bool: whether or not anything was written to the file """ docstr = getdoc(fn) if docstr: reqs = _extract_requirements_from_doc_string(docstr) if reqs: - return _validate_requirements(reqs) + parsed_reqs = _validate_and_parse_requirements(reqs) + _write_requirements_to_file(file_path, parsed_reqs) + return True - return None + return False class FunctionCallsAPI(APIClient): diff --git a/cognite/client/_api/relationships.py b/cognite/client/_api/relationships.py index 22fc25bdec..ca7821c6ad 100644 --- a/cognite/client/_api/relationships.py +++ b/cognite/client/_api/relationships.py @@ -325,7 +325,9 @@ def create( Args: relationship (Union[Relationship, Sequence[Relationship]]): Relationship or list of relationships to create. - Note: the source_type and target_type field in the Relationship(s) can be any string among "Asset", "TimeSeries", "FileMetadata", "Event", "Sequence" + Note: + - the source_type and target_type field in the Relationship(s) can be any string among "Asset", "TimeSeries", "File", "Event", "Sequence"; + - do not provide the value for the source and target arguments of the Relationship class, only source_external_id / source_type and target_external_id / target_type. These (source and target) are used as part of fetching actual resources specified in other fields. Returns: Union[Relationship, RelationshipList]: Created relationship(s) diff --git a/cognite/client/data_classes/functions.py b/cognite/client/data_classes/functions.py index 4a4b4dca7e..080d15118e 100644 --- a/cognite/client/data_classes/functions.py +++ b/cognite/client/data_classes/functions.py @@ -377,7 +377,7 @@ def __init__( self.response_size_mb = response_size_mb @classmethod - def _load(cls, api_response: Dict): # type: ignore # Should be fixed by Self class but it will be available on 3.11 + def _load(cls, api_response: Dict) -> "FunctionsLimits": return cls( timeout_minutes=api_response["timeoutMinutes"], cpu_cores=api_response["cpuCores"], @@ -401,7 +401,7 @@ def __init__( self.status = status @classmethod - def _load(cls, api_response: Dict): # type: ignore # Should be fixed by Self class but it will be available on 3.11 + def _load(cls, api_response: Dict) -> "FunctionsStatus": return cls( status=api_response["status"], ) diff --git a/cognite/client/data_classes/relationships.py b/cognite/client/data_classes/relationships.py index 7c7b1e5df7..147e7d063b 100644 --- a/cognite/client/data_classes/relationships.py +++ b/cognite/client/data_classes/relationships.py @@ -85,7 +85,7 @@ def _validate_resource_types(self) -> "Relationship": @staticmethod def _validate_resource_type(resource_type: Optional[str]) -> None: - _RESOURCE_TYPES = {"asset", "timeseries", "file", "event", "sequence", "geospatial"} + _RESOURCE_TYPES = {"asset", "timeseries", "file", "event", "sequence"} if resource_type is None or resource_type.lower() not in _RESOURCE_TYPES: raise TypeError("Invalid source or target '{}' in relationship".format(resource_type)) diff --git a/tests/tests_unit/test_api/test_functions.py b/tests/tests_unit/test_api/test_functions.py index 1f39e4fb0f..3c368d0243 100644 --- a/tests/tests_unit/test_api/test_functions.py +++ b/tests/tests_unit/test_api/test_functions.py @@ -1,4 +1,5 @@ import os +from tempfile import NamedTemporaryFile from unittest.mock import patch import pytest @@ -8,9 +9,9 @@ from cognite.client._api.functions import ( _extract_requirements_from_doc_string, _extract_requirements_from_file, - _get_requirements_handle, _using_client_credential_flow, - _validate_requirements, + _validate_and_parse_requirements, + _write_fn_docstring_requirements_to_file, validate_function_folder, ) from cognite.client.data_classes import ( @@ -25,7 +26,7 @@ ) from cognite.client.data_classes.functions import FunctionsStatus from cognite.client.exceptions import CogniteAPIError -from tests.utils import jsgz_load +from tests.utils import jsgz_load, set_env_var def post_body_matcher(params): @@ -113,8 +114,8 @@ def match(request_body): @pytest.fixture -def mock_sessions_with_client_credentials(rsps, cognite_client): - url = cognite_client.functions._get_base_url_with_base_path() + "/sessions" +def mock_sessions_with_client_credentials(rsps, cognite_client_with_client_credentials_flow): + url = cognite_client_with_client_credentials_flow.functions._get_base_url_with_base_path() + "/sessions" rsps.add( rsps.POST, @@ -126,8 +127,8 @@ def mock_sessions_with_client_credentials(rsps, cognite_client): { "items": [ { - "clientId": cognite_client.config.token_client_id, - "clientSecret": cognite_client.config.token_client_secret, + "clientId": cognite_client_with_client_credentials_flow.config.token_client_id, + "clientSecret": cognite_client_with_client_credentials_flow.config.token_client_secret, } ] } @@ -305,12 +306,16 @@ def mock_function_calls_filter_response(rsps, cognite_client): yield rsps +@pytest.fixture +def cognite_client_with_client_credentials_flow(): + with set_env_var("COGNITE_API_KEY", "bla"): + client = CogniteClient(token_url="bla", token_scopes=["bla"], token_client_secret="bla", token_client_id="bla") + return client + + @pytest.fixture def cognite_client_with_api_key(): - client = CogniteClient( - api_key="caner_was_here_but_not_for_long_because_api_keys_will_be_removed", - disable_pypi_version_check=True, - ) + client = CogniteClient(api_key="caner_was_here_but_not_for_long_because_api_keys_will_be_removed") client.config.token_client_id = None # Disables Client Credentials coming from the ENV return client @@ -318,12 +323,8 @@ def cognite_client_with_api_key(): @pytest.fixture def cognite_client_with_token(): - client = CogniteClient( - token="aabbccddeeffgg", - disable_pypi_version_check=True, - ) + client = CogniteClient(token="aabbccddeeffgg") client.config.token_client_id = None # Disables Client Credentials coming from the ENV - return client @@ -571,19 +572,21 @@ def test_function_call_timeout_from_api_key_flow( assert mock_functions_call_timeout_response.calls[0].response.json() == res.dump(camel_case=True) @pytest.mark.usefixtures("mock_sessions_with_client_credentials") - def test_function_call_from_oidc_client_credentials_flow(self, mock_functions_call_responses, cognite_client): - assert _using_client_credential_flow(cognite_client) - res = cognite_client.functions.call(id=FUNCTION_ID) + def test_function_call_from_oidc_client_credentials_flow( + self, mock_functions_call_responses, cognite_client_with_client_credentials_flow + ): + assert _using_client_credential_flow(cognite_client_with_client_credentials_flow) + res = cognite_client_with_client_credentials_flow.functions.call(id=FUNCTION_ID) assert isinstance(res, FunctionCall) assert mock_functions_call_responses.calls[2].response.json()["items"][0] == res.dump(camel_case=True) @pytest.mark.usefixtures("mock_sessions_with_client_credentials") def test_function_call_by_external_id_from_oidc_client_credentials_flow( - self, mock_functions_call_by_external_id_responses, cognite_client + self, mock_functions_call_by_external_id_responses, cognite_client_with_client_credentials_flow ): - assert _using_client_credential_flow(cognite_client) - res = cognite_client.functions.call(external_id=f"func-no-{FUNCTION_ID}") + assert _using_client_credential_flow(cognite_client_with_client_credentials_flow) + res = cognite_client_with_client_credentials_flow.functions.call(external_id=f"func-no-{FUNCTION_ID}") assert isinstance(res, FunctionCall) assert mock_functions_call_by_external_id_responses.calls[3].response.json()["items"][0] == res.dump( @@ -591,19 +594,19 @@ def test_function_call_by_external_id_from_oidc_client_credentials_flow( ) @pytest.mark.usefixtures("mock_sessions_bad_request_response") - def test_function_call_with_failing_client_credentials_flow(self, cognite_client): + def test_function_call_with_failing_client_credentials_flow(self, cognite_client_with_client_credentials_flow): with pytest.raises(CogniteAPIError) as excinfo: - assert _using_client_credential_flow(cognite_client) - cognite_client.functions.call(id=FUNCTION_ID) + assert _using_client_credential_flow(cognite_client_with_client_credentials_flow) + cognite_client_with_client_credentials_flow.functions.call(id=FUNCTION_ID) assert "Failed to create session using client credentials flow." in str(excinfo.value) @pytest.mark.usefixtures("mock_sessions_with_client_credentials") def test_function_call_timeout_from_oidc_client_credentials_flow( - self, mock_functions_call_timeout_response, cognite_client + self, mock_functions_call_timeout_response, cognite_client_with_client_credentials_flow ): - assert _using_client_credential_flow(cognite_client) + assert _using_client_credential_flow(cognite_client_with_client_credentials_flow) - res = cognite_client.functions.call(id=FUNCTION_ID) + res = cognite_client_with_client_credentials_flow.functions.call(id=FUNCTION_ID) assert isinstance(res, FunctionCall) assert mock_functions_call_timeout_response.calls[1].response.json() == res.dump(camel_case=True) @@ -668,15 +671,14 @@ class TestRequirementsParser: """Test extraction of requirements.txt from docstring in handle-function""" def test_validate_requirements(self): - out_path = _validate_requirements(["asyncio==3.4.3", "numpy==1.23.0", "pandas==1.4.3"]) - assert os.path.exists(out_path) - os.remove(out_path) + parsed = _validate_and_parse_requirements(["asyncio==3.4.3", "numpy==1.23.0", "pandas==1.4.3"]) + assert parsed == ["asyncio==3.4.3", "numpy==1.23.0", "pandas==1.4.3"] def test_validate_requirements_error(self): reqs = [["asyncio=3.4.3"], ["num py==1.23.0"], ["pandas==1.4.3 python_version=='3.8'"]] for req in reqs: with pytest.raises(Exception): - _validate_requirements(req) + _validate_and_parse_requirements(req) def test_get_requirements_handle(self): def fn(): @@ -687,15 +689,15 @@ def fn(): """ return None - res = _get_requirements_handle(fn=fn) - assert res is not None - os.remove(res) + with NamedTemporaryFile() as ntf: + assert _write_fn_docstring_requirements_to_file(fn, ntf.name) is True def test_get_requirements_handle_error(self): def fn(): return None - assert _get_requirements_handle(fn=fn) is None + with NamedTemporaryFile() as ntf: + assert _write_fn_docstring_requirements_to_file(fn, ntf.name) is False def test_get_requirements_handle_no_docstr(self): def fn(): @@ -707,7 +709,8 @@ def fn(): return None with pytest.raises(Exception): - _get_requirements_handle(fn=fn) is None + with NamedTemporaryFile() as ntf: + assert _write_fn_docstring_requirements_to_file(fn, ntf.name) is False def test_get_requirements_handle_no_reqs(self): def fn(): @@ -717,14 +720,14 @@ def fn(): """ return None - assert _get_requirements_handle(fn=fn) is None + with NamedTemporaryFile() as ntf: + assert _write_fn_docstring_requirements_to_file(fn, ntf.name) is False def test_extract_requirements_from_file(self, tmpdir): req = "somepackage == 3.8.1" file = os.path.join(tmpdir, "requirements.txt") with open(file, "w+") as f: f.writelines("\n".join(["# this should not be included", " " + req])) - f.close() reqs = _extract_requirements_from_file(file_name=file) assert type(reqs) == list assert len(reqs) == 1 @@ -1064,11 +1067,13 @@ def test_get_logs_on_listed_call_object(self, mock_function_call_logs_response, @pytest.mark.usefixtures("mock_sessions_with_client_credentials") @pytest.mark.usefixtures("mock_functions_call_responses") - def test_get_logs_on_created_call_object(self, mock_function_call_logs_response, cognite_client): - call = cognite_client.functions.call(id=FUNCTION_ID) + def test_get_logs_on_created_call_object( + self, mock_function_call_logs_response, cognite_client_with_client_credentials_flow + ): + call = cognite_client_with_client_credentials_flow.functions.call(id=FUNCTION_ID) logs = call.get_logs() assert isinstance(logs, FunctionCallLog) - assert mock_function_call_logs_response.calls[3].response.json()["items"] == logs.dump(camel_case=True) + assert mock_function_call_logs_response.calls[-1].response.json()["items"] == logs.dump(camel_case=True) def test_function_call_response_by_function_id(self, mock_function_call_response_response, cognite_client): res = cognite_client.functions.calls.get_response(call_id=CALL_ID, function_id=FUNCTION_ID) diff --git a/tests/tests_unit/test_data_classes/test_contextualization.py b/tests/tests_unit/test_data_classes/test_contextualization.py index 744f66156d..f42200e9e9 100644 --- a/tests/tests_unit/test_data_classes/test_contextualization.py +++ b/tests/tests_unit/test_data_classes/test_contextualization.py @@ -2,15 +2,14 @@ import pytest -from cognite.client import CogniteClient from cognite.client.data_classes import ContextualizationJob from tests.utils import set_env_var @pytest.fixture() -def job(): +def job(cognite_client): with set_env_var("COGNITE_API_KEY", "BLA"): - return ContextualizationJob(job_id=123, status="Queued", cognite_client=CogniteClient()) + return ContextualizationJob(job_id=123, status="Queued", cognite_client=cognite_client) def mock_update_status_running(self):