Python 3.9：从 Cloud Datastore 检索数据的单元测试

Question

我是使用 python 进行单元测试的新手，我在尝试为我的代码编写一个代码时遇到了麻烦，我在其中获取了两个 Datastore entities 和根据条件设置时间戳。

我的main.py 看起来像这样：

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

所以我从 Datastore 中获取两个实体及其属性，我根据条件选择时间戳属性，然后比较两个时间戳。

通过搜索，我只知道我可以使用 patch 来模拟 Datastore API，test_main.py 看起来像这样：

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

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

谁能帮我告诉我代码的单元测试是什么样子的？ 如果有帮助，我正在使用 Python 3.9。

感谢您的帮助，并提前致谢。

Answer 1

我现在设法编写了单元测试。

首先，我将函数拆分为 2 个较小的函数，以便测试代码的较小部分：

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

我能够编写这样的测试：

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"

对于实体，我从 json 文件中获取了一些示例数据。