当前位置：首页 > article >正文

Autogen_core 测试代码：test_cache_store.py

article 2025/1/30 10:07:17

原始代码

from typing import Dict, Generic, Optional, Protocol, TypeVar

T = TypeVar("T")


class CacheStore(Protocol, Generic[T]):
    """
    This protocol defines the basic interface for store/cache operations.

    Sub-classes should handle the lifecycle of underlying storage.
    """

    def get(self, key: str, default: Optional[T] = None) -> Optional[T]:
        """
        Retrieve an item from the store.

        Args:
            key: The key identifying the item in the store.
            default (optional): The default value to return if the key is not found.
                                Defaults to None.

        Returns:
            The value associated with the key if found, else the default value.
        """
        ...

    def set(self, key: str, value: T) -> None:
        """
        Set an item in the store.

        Args:
            key: The key under which the item is to be stored.
            value: The value to be stored in the store.
        """
        ...


class InMemoryStore(CacheStore[T]):
    def __init__(self) -> None:
        self.store: Dict[str, T] = {}

    def get(self, key: str, default: Optional[T] = None) -> Optional[T]:
        return self.store.get(key, default)

    def set(self, key: str, value: T) -> None:
        self.store[key] = value

这段代码定义了一个缓存存储的基本接口和一个基于内存的实现。

首先，它从typing模块导入了一些泛型类型定义，用于后续的类型注解。
接着，定义了一个类型变量T，这是一个泛型占位符，用于表示存储在缓存中的数据的类型。
然后，定义了一个名为CacheStore的协议类，它是一个泛型类，参数为T。这个类定义了缓存存储操作的基本接口，包括get和set方法。get方法用于从存储中检索一个项，如果找不到指定的键，则返回默认值；set方法用于在存储中设置一个项。
最后，定义了一个名为InMemoryStore的类，它继承自CacheStore并实现了get和set方法。InMemoryStore使用一个字典self.store来在内存中存储数据。get方法通过键来获取存储中的值，如果键不存在，则返回默认值；set方法用于将一个键值对存储到字典中。

总的来说，这段代码定义了一个缓存存储的接口和一个简单的内存存储实现，使得可以通过键来存取数据，并提供了一个默认值机制。

测试代码

from unittest.mock import Mock

from autogen_core import CacheStore, InMemoryStore


def test_set_and_get_object_key_value() -> None:
    mock_store = Mock(spec=CacheStore)
    test_key = "test_key"
    test_value = object()
    mock_store.set(test_key, test_value)
    mock_store.get.return_value = test_value
    mock_store.set.assert_called_with(test_key, test_value)
    assert mock_store.get(test_key) == test_value

test_set_and_get_object_key_value()

第一段代码定义了一个名为 test_set_and_get_object_key_value 的测试函数。这个函数测试了一个缓存存储对象（通过 Mock 模拟的）的 set 和 get 方法。具体步骤如下：

创建一个 Mock 对象 mock_store，它模拟了 CacheStore 类。
定义了一个测试用的键 test_key 和一个测试用的值 test_value（这里是一个新的 object 实例）。
调用 mock_store.set(test_key, test_value) 来模拟设置键值对。
设置 mock_store.get 方法的返回值为 test_value，这样当调用 get 方法时，它会返回之前设置的值。
通过 mock_store.set.assert_called_with(test_key, test_value) 断言 set 方法是否被正确调用。
断言 mock_store.get(test_key) 的返回值是否等于 test_value。

这个测试函数检查了存储对象是否能够正确地设置和获取键值对。

def test_get_non_existent_key() -> None:
    mock_store = Mock(spec=CacheStore)
    key = "non_existent_key"
    mock_store.get.return_value = None
    assert mock_store.get(key) is None


test_get_non_existent_key()

第二段代码定义了一个名为 test_get_non_existent_key 的测试函数。这个函数测试了当请求一个不存在的键时，缓存存储对象的行为。具体步骤如下：

创建一个 Mock 对象 mock_store，模拟了 CacheStore 类。
定义一个不存在的键 key。
设置 mock_store.get 方法的返回值为 None。
断言 mock_store.get(key) 的返回值是否为 None。

这个测试函数检查了当请求的键不存在时，存储对象是否返回 None。

def test_set_overwrite_existing_key() -> None:
    mock_store = Mock(spec=CacheStore)
    key = "test_key"
    initial_value = "initial_value"
    new_value = "new_value"
    mock_store.set(key, initial_value)
    mock_store.set(key, new_value)
    mock_store.get.return_value = new_value
    mock_store.set.assert_called_with(key, new_value)
    assert mock_store.get(key) == new_value

test_set_overwrite_existing_key()

第三段代码定义了一个名为 test_set_overwrite_existing_key 的测试函数。这个函数测试了设置一个已经存在的键时，缓存存储对象是否会覆盖旧的值。具体步骤如下：

创建一个 Mock 对象 mock_store，模拟了 CacheStore 类。
定义了一个键 key 和两个值 initial_value、new_value。
调用 mock_store.set(key, initial_value) 来模拟设置初始键值对。
再次调用 mock_store.set(key, new_value) 来模拟覆盖设置键值对。
设置 mock_store.get 方法的返回值为 new_value。
断言 mock_store.set 是否被用 key 和 new_value 调用。
断言 mock_store.get(key) 的返回值是否等于 new_value。

这个测试函数检查了设置已存在的键时，存储对象是否正确地覆盖了旧的值。

def test_inmemory_store() -> None:
    store = InMemoryStore[int]()
    test_key = "test_key"
    test_value = 42
    store.set(test_key, test_value)
    assert store.get(test_key) == test_value

    new_value = 2
    store.set(test_key, new_value)
    assert store.get(test_key) == new_value

    key = "non_existent_key"
    default_value = 99
    assert store.get(key, default_value) == default_value
    
test_inmemory_store()