Python 3.9: Unit testing of retrieving data from Cloud Datastore

user_14132084 0
Viewed 208

I'm new to unit tests with python and I have trouble trying to write one for my code where I fetch a few properties (string, timestamp) of two Datastore entities and based on a condition set the timestamp.

My main.py looks like this:

from google.cloud import datastore

datastore_client = datastore.Client()

def get_timestamp():
    # fetch property1 (string), property2 (timestamp), property3 (timestamp) from entity1
    query = datastore_client.query(kind="kind1")
    key_entity1 = datastore_client.key("kind1", "entity1_key_id")
    query.key_filter(key_entity1, "=")
    list_entity1 = list(query.fetch())
    entity1 = dict(list_entity1[0])
    property1 = entity1['property1']
    # based on the value of property1 set timestamp1
    if property1 == "value1" or property1 == "value2":
        timestamp1 = entity1['property2']
    elif property1 == "value3":
        timestamp1 = entity1['property3']
    # fetch property1 (timestamp) from entity2
    query = datastore_client.query(kind="kind1")
    key_entity2 = datastore_client.key("job_results", "entity2_key_id")
    query.key_filter(key_entity2, "=")
    list_entity2 = list(query.fetch())
    entity2 = dict(list_entity2[0])
    timestamp2 = entity2['property1']
    if timestamp2 > timestamp1:
        timestamp = timestamp2.isoformat().replace('+00:00', 'Z')
    elif timestamp1 > timestamp2:
        timestamp = timestamp1.isoformat().replace('+00:00', 'Z')
    return timestamp

So I'm fetching two entities and their properties from Datastore, I choose the timestamp property based on a condition and then I compare two timestamps.

From searching I just know I can use patch to mock the Datastore API and test_main.py can look like this:

import pytest
from unittest.mock import Mock, patch
import main

@patch("main.datastore_client")
def test_get_timestamp():
    # test code

Can someone help and tell me how the unit test looks like for my code? If it helps I'm using Python 3.9.

I appreciate any help and thanks in advance.

google-api google-cloud-datastore unit-testing python
1 Answers

I managed to write the unit tests now.

First I split the function into 2 smaller ones so I can test smaller parts of the code:

from google.cloud import datastore

datastore_client = datastore.Client()

def get_entity(datastore_client, entity):
    query = datastore_client.query(kind="kind1")
    key_entity = datastore_client.key("kind1", entity)
    query.key_filter(key_entity, "=")
    list_entity = list(query.fetch())
    return list_entity

list_entity1 = get_entity(datastore_client, "entity1_key_id")
            list_entity2 = get_entity(datastore_client, "entity2_key_id")

def get_timestamp(list_entity1, list_entity2):
    entity1 = dict(list_entity1[0])
    property1 = entity1['property1']
    # based on the value of property1 set timestamp1
    if property1 == "value1" or property1 == "value2":
        timestamp1 = entity1['property2']
    elif property1 == "value3":
        timestamp1 = entity1['property3']
    entity2 = dict(list_entity2[0])
    timestamp2 = entity2['property1']
    if timestamp2 > timestamp1:
        timestamp = timestamp2.isoformat().replace('+00:00', 'Z')
    elif timestamp1 > timestamp2:
        timestamp = timestamp1.isoformat().replace('+00:00', 'Z')
    return timestamp

I was able to write the tests like that:

import pytest
from unittest.mock import MagicMock
import main as test

@pytest.fixture(scope="function")
def mock_query():
    mock_query = MagicMock()
    mock_query.key_filter = MagicMock()
    mock_query.fetch = MagicMock(return_value=[{}])
    return mock_query

@pytest.fixture(scope="function")
def mock_client(mock_query):
    mock_client = MagicMock()
    mock_client.query = MagicMock(return_value=mock_query)
    mock_client.key = MagicMock()
    return mock_client

def test_get_entity(mock_query: MagicMock, mock_client: MagicMock):
    mock_entity  = MagicMock()
    list_entity = test.get_entity(mock_client, mock_entity)
    mock_client.query.assert_called_once()
    mock_client.key.assert_called_once()
    mock_query.key_filter.assert_called_once_with(mock_client.key.return_value, "=")
    mock_query.fetch.assert_called_once()
    assert list_entity == mock_query.fetch.return_value

# covers condition if timestamp2 is the latest
def test_get_timestamp():
    entity1 = read_json_to_dict_entity1("test_data_entity1.json")
    list_entity1 = [entity1]
    entity2 = read_json_to_dict_entity2("test_data_entity2.json")
    list_entity2 = [entity2]
    timestamp = test.get_timestamp(list_entity1, list_entity2)
    assert timestamp == "timestamp2"

For the entities I took some example data from json files.

user_141320840
Related