Add context manager for mocking CogniteClient (#503)

CHANGELOG.md

@@ -25,6 +25,7 @@ Changes are grouped as follows
 - SequencesAPI.list now accepts an asset_ids parameter for searching by asset
 - SequencesDataAPI.insert now accepts a SequenceData object for easier copying
 - DatapointsAPI.insert now accepts a Datapoints object for easier copying
+- helper method `cognite.client.testing.mock_cognite_client()` for mocking CogniteClient
 
 ### Changed
 - assets.create() no longer validates asset hierarchy and sorts assets before posting. This functionality has been moved to assets.create_hierarchy().

Pipfile

@@ -7,6 +7,7 @@ name = "pypi"
 pandas = "*"
 requests = "*"
 matplotlib = "*"
+pytest = "<5"
 
 [dev-packages]
 sphinx = "*"

cognite/client/data_classes/__init__.py

@@ -14,6 +14,7 @@
     ServiceAccount,
     ServiceAccountList,
 )
+from cognite.client.data_classes.login import LoginStatus
 from cognite.client.data_classes.raw import Database, DatabaseList, Row, RowList, Table, TableList
 from cognite.client.data_classes.sequences import Sequence, SequenceData, SequenceFilter, SequenceList, SequenceUpdate
 from cognite.client.data_classes.three_d import (

cognite/client/data_classes/login.py

@@ -7,13 +7,14 @@ class LoginStatus(CogniteResponse):
     """Current login status
 
     Args:
-        user (str): Current user
-        logged_in (bool): Is user logged in
-        project (str): Current project
-        project_id (str): Current project id
+        user (str): Current user.
+        logged_in (bool): Is user logged in.
+        project (str): Current project.
+        project_id (int): Current project id.
+        api_key_id (int): Current api key id.
     """
 
-    def __init__(self, user: str, project: str, logged_in: bool, project_id: str, api_key_id: int):
+    def __init__(self, user: str, project: str, logged_in: bool, project_id: int, api_key_id: int):
         self.user = user
         self.project = project
         self.project_id = project_id

cognite/client/testing.py

@@ -0,0 +1,96 @@
+from contextlib import contextmanager
+from unittest import mock
+
+from cognite.client import CogniteClient
+from cognite.client._api.assets import AssetsAPI
+from cognite.client._api.datapoints import DatapointsAPI
+from cognite.client._api.events import EventsAPI
+from cognite.client._api.files import FilesAPI
+from cognite.client._api.iam import IAMAPI, APIKeysAPI, GroupsAPI, SecurityCategoriesAPI, ServiceAccountsAPI
+from cognite.client._api.login import LoginAPI
+from cognite.client._api.raw import RawAPI, RawDatabasesAPI, RawRowsAPI, RawTablesAPI
+from cognite.client._api.three_d import (
+    ThreeDAPI,
+    ThreeDAssetMappingAPI,
+    ThreeDFilesAPI,
+    ThreeDModelsAPI,
+    ThreeDRevisionsAPI,
+)
+from cognite.client._api.time_series import TimeSeriesAPI
+
+
+@contextmanager
+def mock_cognite_client():
+    """Context manager for mocking the CogniteClient.
+
+    Will patch all APIs and replace the client with specced MagicMock objects.
+
+    Examples:
+
+        In this example we can run the following code without actually executing the underlying API calls::
+
+            >>> from cognite.client import CogniteClient
+            >>> from cognite.client.data_classes import TimeSeries
+            >>> from cognite.client.testing import mock_cognite_client
+            >>>
+            >>> with mock_cognite_client():
+            >>>     c = CogniteClient()
+            >>>     c.time_series.create(TimeSeries(external_id="blabla"))
+
+        This example shows how to set the return value of a given method::
+
+            >>> from cognite.client import CogniteClient
+            >>> from cognite.client.data_classes import TimeSeries
+            >>> from cognite.client.data_classes import LoginStatus
+            >>> from cognite.client.testing import mock_cognite_client
+            >>>
+            >>> with mock_cognite_client() as c_mock:
+            >>>     c_mock.login.status.return_value = LoginStatus(
+            >>>         user="user", project="dummy", project_id=1, logged_in=True, api_key_id=1
+            >>>     )
+            >>>     c = CogniteClient()
+            >>>     res = c.login.status()
+            >>>     assert "user" == res.user
+
+        Here you can see how to have a given method raise an exception::
+
+            >>> from cognite.client import CogniteClient
+            >>> from cognite.client.exceptions import CogniteAPIError
+            >>> from cognite.client.testing import mock_cognite_client
+            >>>
+            >>> with mock_cognite_client() as c_mock:
+            >>>     c_mock.login.status.side_effect = CogniteAPIError(message="Something went wrong", code=400)
+            >>>     c = CogniteClient()
+            >>>     try:
+            >>>         res = c.login.status()
+            >>>     except CogniteAPIError as e:
+            >>>         assert 400 == e.code
+            >>>         assert "Something went wrong" == e.message
+    """
+    client_new_fn = CogniteClient.__new__
+
+    cog_client_mock = mock.MagicMock(spec=CogniteClient)
+    cog_client_mock.time_series = mock.MagicMock(spec_set=TimeSeriesAPI)
+    cog_client_mock.datapoints = mock.MagicMock(spec_set=DatapointsAPI)
+    cog_client_mock.assets = mock.MagicMock(spec_set=AssetsAPI)
+    cog_client_mock.events = mock.MagicMock(spec_set=EventsAPI)
+    cog_client_mock.files = mock.MagicMock(spec_set=FilesAPI)
+    cog_client_mock.login = mock.MagicMock(spec_set=LoginAPI)
+    cog_client_mock.three_d = mock.MagicMock(spec=ThreeDAPI)
+    cog_client_mock.three_d.models = mock.MagicMock(spec_set=ThreeDModelsAPI)
+    cog_client_mock.three_d.revisions = mock.MagicMock(spec_set=ThreeDRevisionsAPI)
+    cog_client_mock.three_d.files = mock.MagicMock(spec_set=ThreeDFilesAPI)
+    cog_client_mock.three_d.asset_mappings = mock.MagicMock(spec_set=ThreeDAssetMappingAPI)
+    cog_client_mock.iam = mock.MagicMock(spec=IAMAPI)
+    cog_client_mock.iam.service_accounts = mock.MagicMock(spec=ServiceAccountsAPI)
+    cog_client_mock.iam.api_keys = mock.MagicMock(spec_set=APIKeysAPI)
+    cog_client_mock.iam.groups = mock.MagicMock(spec_set=GroupsAPI)
+    cog_client_mock.iam.security_categories = mock.MagicMock(spec_set=SecurityCategoriesAPI)
+    cog_client_mock.raw = mock.MagicMock(spec=RawAPI)
+    cog_client_mock.raw.databases = mock.MagicMock(spec_set=RawDatabasesAPI)
+    cog_client_mock.raw.tables = mock.MagicMock(spec_set=RawTablesAPI)
+    cog_client_mock.raw.rows = mock.MagicMock(spec_set=RawRowsAPI)
+
+    CogniteClient.__new__ = lambda x: cog_client_mock
+    yield cog_client_mock
+    CogniteClient.__new__ = client_new_fn

docs/source/cognite.rst

@@ -779,6 +779,11 @@ Convert milliseconds since epoch to datetime
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 .. autofunction:: cognite.client.utils.ms_to_datetime
 
+Testing
+-------
+Use a context manager to mock CogniteClient
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. autofunction:: cognite.client.testing.mock_cognite_client
 
 Experimental features
 =====================

tests/tests_unit/test_testing.py

@@ -0,0 +1,31 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from cognite.client import CogniteClient
+from cognite.client.testing import mock_cognite_client
+
+
+def test_mock_cognite_client():
+    with mock_cognite_client() as c_mock:
+        c = CogniteClient()
+        assert isinstance(c_mock, MagicMock)
+        assert c_mock == c
+
+        api_pairs = [
+            (c.time_series, c_mock.time_series),
+            (c.raw, c_mock.raw),
+            (c.assets, c_mock.assets),
+            (c.datapoints, c_mock.datapoints),
+            (c.events, c_mock.events),
+            (c.files, c_mock.files),
+            (c.iam, c_mock.iam),
+            (c.login, c_mock.login),
+            (c.three_d, c_mock.three_d),
+        ]
+        for api, mock_api in api_pairs:
+            assert isinstance(mock_api, MagicMock)
+            assert api == mock_api
+
+            with pytest.raises(AttributeError):
+                api.does_not_exist