How to test Python code involving SDK client calls [Cognite Official]

  • 15 December 2022
  • 0 replies
  • 301 views

Userlevel 5

Automated tests are considered a good practice in software development. They can apply to every language or use case. Here we will focus on Python and testing code that involves the Python SDK’s CogniteClient object. Below are listed some of the benefits of testing. 

 

First, they are a good way to test your code in a fast and repeated manner. With automated tests, you avoid testing manually your code: you can write each test case programmatically and run them on request. When they fail, running them in a debugger really helps understand the code, each object’s structure, and solve the problem quite fast.

 

Then, it helps make the code maintainable. Indeed, when the code needs to be changed, running your tests ensures that the tested code will still work as expected. When a bug is discovered, it is good to create a test case for that bug to avoid having it again when updating the code.

 

It also helps collaborate on the same repository. When several people work on the same code, it might happen that someone writes something that has breaking changes in someone else’s code. Automated tests help prevent this: tests will fail when run if there are such breaking changes.


 

Testing logic that involves calls to the Cognite API, with the Python SDK, can look a bit tricky at first. Indeed, we do not make actual calls to the API during testing: we do not want to have an authenticated client for that, that would slow down the tests, and the tests would also rely on things like the network. We will use a mock client instead. Also, we do not want to test the API and the SDK themselves: they are already tested. What we want to test is the logic that happens before and after a client call. Since the client call will still happen, no matter what, in the logic, we will replace it by a mock client for the tests.

 

A mock client is not an actual client, but it still has the same endpoints, for which you can programmatically set return values and side effects (we’ll have a further look later). The Cognite SDK provides a mock client: monkeypatch_cognite_client (read the docs here: https://cognite-sdk-python.readthedocs-hosted.com/en/latest/cognite.html#testing). 


 

A good practice for using a client in your functions is using dependency injection. Basically, that means we should pass our client as a function argument, instead of instantiating it in the function code. Feel free to read more about that topic. One of the benefits is that it makes it easier to test: we can then pass a mock client instead of an actual instantiated client.

 

Let’s say we want to create two functions: 

  • The first one should take as input a list of dictionaries (with external_id, name, some_number as arguments) and a client, and create assets in CDF with external_id as external_id, name as name, and two metadata fields: some_number and some_number_times_2 (2 * some_number). It does not return anything. In case the list is empty, the client will raise a ValueError exception (nothing to change here).

  • The second takes as input a list of external_ids (strings) and a client. It first filters the list to keep only the external_ids starting with “a”. Then it makes a call to CDF to retrieve the corresponding assets. Finally it returns a list of those assets’ names. In case the input list is empty, we want to return an empty list, with no exception raised.

 

In theory, you could define your tests before actually writing your function. For the purpose of making the article more understandable, you will already find below the two functions: 


def create_assets(items: List[Dict], client: CogniteClient):

assets = [

Asset(

external_id=item["external_id"],

name=item["name"],

metadata={

"some_number": str(item["some_number"]),

"some_number_times_2": str(item["some_number"] * 2),

},

)

for item in items

]

client.assets.create(assets)




def retrieve_name_of_assets_starting_with_a_from_list(

items: List[str], client: CogniteClient

):

asset_external_ids = [item for item in items if item.startswith("a")]

try:

assets = client.assets.retrieve_multiple(external_ids=asset_external_ids)

except ValueError:

assets = []

return [asset.name for asset in assets]

 

First, let’s see how to set up such a mock client. We’ll also use pytest (and its fixtures) for testing. We want to setup return values for the client.assets.create and client.assets.retrieve_multiple values.

CREATED_ASSETS = [

Asset(

external_id="abc",

name="abc",

metadata={"some_number": "2", "some_number_times_2": "4"},

),

Asset(

external_id="abcd",

name="abcd",

metadata={"some_number": "3", "some_number_times_2": "6"},

),

Asset(

external_id="bcd",

name="bcd",

metadata={"some_number": "2", "some_number_times_2": "4"},

),

Asset(

external_id="bcde",

name="bcde",

metadata={"some_number": "3", "some_number_times_2": "6"},

),

]



RETRIEVED_ASSETS_EXTERNAL_ID_STARTS_WITH_A = [

Asset(

external_id="abc",

name="abc",

metadata={"some_number": "2", "some_number_times_2": "4"},

),

Asset(

external_id="abcd",

name="abcd",

metadata={"some_number": "3", "some_number_times_2": "6"},

),

]




@pytest.fixture

def cognite_client_mock():

with monkeypatch_cognite_client() as client:

client.assets.create.return_value = CREATED_ASSETS

client.assets.retrieve_multiple.return_value = RETRIEVED_ASSETS_EXTERNAL_ID_STARTS_WITH_A

return client

 

The role of the pytest fixture is to make the decorated function usable as a variable in the tests. What this setup will do is that, for example, every time the client.assets.create endpoint is called, it will return the CREATED_ASSETS list.

 

The test for the first function is the easiest to start with. Since it ends with a client call and does not return anything, we simply need to check that the client is called with the correct arguments. The test looks like this: 

def test_create_assets(cognite_client_mock):

items = [

{"external_id": "abc", "name": "abc", "some_number": 2},

{"external_id": "abcd", "name": "abcd", "some_number": 3},

{"external_id": "bcd", "name": "bcd", "some_number": 2},

{"external_id": "bcde", "name": "bcde", "some_number": 3},

]

create_assets(items, cognite_client_mock)

assert cognite_client_mock.assets.create.call_count == 1

assert len(cognite_client_mock.assets.create.call_args.args[0]) == 4

assert [

asset.dump() for asset in cognite_client_mock.assets.create.call_args.args[0]

] == [asset.dump() for asset in CREATED_ASSETS]

 

In that test, we check that the client.assets.create endpoint is called only once. We also check that it is called with the expected arguments (positional arguments here, or args).


 

The test for the second function is a bit more complex: we need to check that the client is called with the correct arguments, but also that what the client returns is processed as expected. Also, we want to check that exceptions are handled correctly: by setting an exception as a side effect, the mock client will return a ValueError when the assets.retrieve_multiple endpoint is called.

def test_retrieve_name_of_assets_starting_with_a_from_list(cognite_client_mock):

external_ids = ["abc", "abcd", "bcd", "bcde"]

names = retrieve_name_of_assets_starting_with_a_from_list(

external_ids, cognite_client_mock

)

assert cognite_client_mock.assets.retrieve_multiple.call_count == 1

assert cognite_client_mock.assets.retrieve_multiple.call_args.kwargs == {

"external_ids": ["abc", "abcd"]

}

assert names == ["abc", "abcd"]



external_ids = ["bcd", "bcde"]

cognite_client_mock.assets.retrieve_multiple.side_effect = ValueError()

names = retrieve_name_of_assets_starting_with_a_from_list(

external_ids, cognite_client_mock

)

assert cognite_client_mock.assets.retrieve_multiple.call_count == 2

assert cognite_client_mock.assets.retrieve_multiple.call_args.kwargs == {

"external_ids": []

}

assert names == []

 

In that test, we check that the client.assets.create endpoint is called only once. We also check that it is called with the expected arguments (keyword arguments here, or kwargs). That test also checks that what happens after the client call is correct, based on the return_value we set previously. For the second part of that test, we also set a side effect to the called endpoint, to raise an Exception, in order to assess that the function behaves as expected when this happens.

 

Through this article we’ve gone through some examples of testing functions that involve Cognite client calls. This is made easier thanks to a client mock which allows testing the logic before and after the client call, considering the Cognite SDK itself is already tested. That way we don’t need an instantiated client, with credentials etc. Testing code is a good practice to make it maintainable and make collaboration easier. 

 

@Sigurd Hylin , I know you are quite experienced with testing, any best practice you would like to share ? 

 


0 replies

Be the first to reply!

Reply