From 4210f2f3cf9efb257467a8fa45bef25d9932457a Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Fri, 18 Jul 2025 09:20:20 -0700 Subject: [PATCH 01/13] Adding Cosmos DB storage --- .../microsoft/agents/cosmos/__init__.py | 7 + .../microsoft/agents/cosmos/cosmos.py | 177 ++++++++ .../agents/cosmos/cosmos_db_storage.py | 161 +++++++ .../agents/cosmos/cosmos_db_storage_config.py | 88 ++++ .../microsoft/agents/cosmos/key_ops.py | 37 ++ .../microsoft/agents/cosmos/test.py | 419 ++++++++++++++++++ 6 files changed, 889 insertions(+) create mode 100644 libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/__init__.py create mode 100644 libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos.py create mode 100644 libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py create mode 100644 libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py create mode 100644 libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py create mode 100644 libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/test.py diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/__init__.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/__init__.py new file mode 100644 index 00000000..00d54854 --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/__init__.py @@ -0,0 +1,7 @@ +from .cosmos_db_storage import CosmosDBStorage +from .cosmos_db_storage_config import CosmosDBStorageConfig + +__all__ = [ + "CosmosDBStorage", + "CosmosDBStorageConfig", +] diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos.py new file mode 100644 index 00000000..236be502 --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos.py @@ -0,0 +1,177 @@ +# # Copyright (c) Microsoft Corporation. All rights reserved. +# # Licensed under the MIT License. + +# import json + +# from typing import TypeVar +# from threading import Lock +# from hashlib import sha256 + +# from azure.core.auth import ( +# StorageSharedKeyCredential, +# TokenCredential, +# AnonymousCredential, +# ) + +# from azure.cosmos import ( +# documents, +# http_constants, +# CosmosDict, +# ) +# from azure.cosmos.aio import ( +# ContainerProxy, +# CosmosClient, +# DatabaseProxy, +# ) +# import azure.cosmos.exceptions as cosmos_exceptions + +# from microsoft.agents.storage import Storage, StoreItem +# from microsoft.agents.storage._type_aliases import JSON +# from microsoft.agents.storage.error_handling import ignore_error, is_status_code_error + +# StoreItemT = TypeVar("StoreItemT", bound=StoreItem) + +# class CosmosDBStorage(Storage): +# """A CosmosDB based storage provider using partitioning""" + +# def __init__(self, config: CosmosDBConfig): +# """Create the storage object. + +# :param config: +# """ +# super().__init__() + +# CosmosDBStorage._validate_config(config) + +# self._config: CosmosDBConfig = config +# self._client: CosmosClient = self._create_client() +# self._database: DatabaseProxy = None +# self._container: ContainerProxy = None +# self._compatability_mode_partition_key: bool = False +# # Lock used for synchronizing container creation +# self._lock: Lock = Lock() + +# @staticmethod +# def _validate_config(config: CosmosDBConfig): +# if not config: +# raise ValueError("CosmosDBStorage: CosmosDBConfig is required.") +# if not config.cosmos_db_endpoint: +# raise ValueError("CosmosDBStorage: cosmos_db_endpoint is required.") +# if not config.auth_key: +# raise ValueError("CosmosDBStorage: auth_key is required.") +# if not config.database_id: +# raise ValueError("CosmosDBStorage: database_id is required.") +# if not config.container_id: +# raise ValueError("CosmosDBStorage: container_id is required.") + +# @staticmethod +# def _validate_suffix(config: CosmosDBConfig): + +# if config.key_suffix: +# if config.compatibility_mode: +# raise ValueError( +# "compatibilityMode cannot be true while using a keySuffix." +# ) +# suffix_escaped: str = sanitize_key(config.key_suffix) +# if suffix_escaped != config.key_suffix: +# raise ValueError( +# f"Cannot use invalid Row Key characters: {config.key_suffix} in keySuffix." +# ) + +# def _create_client(self) -> CosmosClient: + +# if self._config.url: +# if not self._config.credential: +# raise ValueError( +# "CosmosDBStorage: Credential is required when using a custom service URL." +# ) +# return CosmosClient(account_url=self._config.url, credential=self._config.credential) +# else: + +# connection_policy = self._config.cosmos_client_options.get( +# "connection_policy", documents.ConnectionPolicy() +# ) + +# # kwargs 'connection_verify' is to handle CosmosClient overwriting the +# # ConnectionPolicy.DisableSSLVerification value. +# return CosmosClient( +# self._config.cosmos_db_endpoint, +# self._config.auth_key, +# self._config.cosmos_client_options.get("consistency_level", None), +# **{ +# "connection_policy": connection_policy, +# "connection_verify": not connection_policy.DisableSSLVerification, +# }, +# ) + +# async def _read_item(key: str) -> tuple[str, StoreItemT]: +# escaped_key: str = sanitize_key( +# key, self._config.key_suffix, self._config.compatibility_mode +# ) +# read_item_response: CosmosDict = await self._container.read_item( +# escaped_key, self._get_partition_key(escaped_key) +# ) +# doc: JSON = read_item_response.get("document") +# return read_item_response["realId"], target_cls.from_json_to_store_item(doc) + +# async def _write_item(key: str, item: StoreItem) -> None: +# doc = { +# "id": sanitize_key( +# key, self._config.key_suffix, self._config.compatibility_mode +# ), +# "realId": key, +# "document": item.store_item_to_json(), +# } +# await self._container.upsert_item(body=doc) + +# async def _delete_item(self, key: str) -> None: +# escaped_key: str = sanitize_key( +# key, self._config.key_suffix, self._config.compatibility_mode +# ) +# await ignore_error(self._container.delete_item( +# escaped_key, self._get_partition_key(escaped_key) +# ), is_status_code_error(404)) + +# async def initialize(self) -> None: +# if not self._container: +# with self._lock: +# # in case another thread attempted to initialize just before acquiring the lock +# if self._container: return + +# if not self._database: +# self._database = self._client.create_database_if_not_exists( +# self._config.database_id +# ) + +# partition_key = { +# "paths": ["/id"], +# "kind": documents.PartitionKind.Hash, +# } +# try: +# self._container = self._database.create_container( +# self._config.container_id, +# partition_key, +# offer_throughput=self._config.container_throughput, +# ) +# except cosmos_exceptions.CosmosHttpResponseError as err: +# if err.status_code == http_constants.StatusCodes.CONFLICT: +# self._container = self._database.get_container_client( +# self._config.container_id +# ) +# properties = self._container.read() +# if "partitionKey" not in properties: +# self._compatability_mode_partition_key = True +# else: +# paths = properties["partitionKey"]["paths"] +# if "/partitionKey" in paths: +# self._compatability_mode_partition_key = True +# elif "/id" not in paths: +# raise Exception( +# f"Custom Partition Key Paths are not supported. {self._config.container_id} " +# "has a custom Partition Key Path of {paths[0]}." +# ) +# else: +# raise err + +# def _get_partition_key(self, key: str) -> str: +# return "" if self._compatability_mode_partition_key else key diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py new file mode 100644 index 00000000..8932e9c2 --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py @@ -0,0 +1,161 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +from typing import TypeVar +from threading import Lock + +from azure.cosmos import ( + documents, + http_constants, + CosmosDict, +) +from azure.cosmos.aio import ( + ContainerProxy, + CosmosClient, + DatabaseProxy, +) +import azure.cosmos.exceptions as cosmos_exceptions + +from microsoft.agents.storage import Storage, StoreItem +from microsoft.agents.storage._type_aliases import JSON +from microsoft.agents.storage.error_handling import ignore_error, is_status_code_error + +from .cosmos_db_storage_config import CosmosDBStorageConfig +from .key_ops import sanitize_key + +StoreItemT = TypeVar("StoreItemT", bound=StoreItem) + +class CosmosDBStorage(Storage): + """A CosmosDB based storage provider using partitioning""" + + def __init__(self, config: CosmosDBStorageConfig): + """Create the storage object. + + :param config: + """ + super().__init__() + + CosmosDBStorageConfig.validate_cosmos_db_config(config) + + self._config: CosmosDBStorageConfig = config + self._client: CosmosClient = self._create_client() + self._database: DatabaseProxy = None + self._container: ContainerProxy = None + self._compatability_mode_partition_key: bool = False + # Lock used for synchronizing container creation + self._lock: Lock = Lock() + + def _create_client(self) -> CosmosClient: + + if self._config.url: + if not self._config.credential: + raise ValueError( + "CosmosDBStorage: Credential is required when using a custom service URL." + ) + return CosmosClient(account_url=self._config.url, credential=self._config.credential) + + connection_policy = self._config.cosmos_client_options.get( + "connection_policy", documents.ConnectionPolicy() + ) + + # kwargs 'connection_verify' is to handle CosmosClient overwriting the + # ConnectionPolicy.DisableSSLVerification value. + return CosmosClient( + self._config.cosmos_db_endpoint, + self._config.auth_key, + consistency_level=self._config.cosmos_client_options.get("consistency_level", None), + **{ + "connection_policy": connection_policy, + "connection_verify": not connection_policy.DisableSSLVerification, + }, + ) + + async def _read_item( + self, key: str, *, target_cls: StoreItemT = None, **kwargs + ) -> tuple[str | None, StoreItemT | None]: + + if key == "": + raise ValueError("CosmosDBStorage: Key cannot be empty.") + + escaped_key: str = sanitize_key( + key, self._config.key_suffix, self._config.compatibility_mode + ) + read_item_response: CosmosDict = await ignore_error(self._container.read_item( + escaped_key, self._get_partition_key(escaped_key) + ), lambda err: isinstance(err, cosmos_exceptions.CosmosResourceNotFoundError)) + + if read_item_response is None: + return None, None + + doc: JSON = read_item_response.get("document") + return read_item_response["realId"], target_cls.from_json_to_store_item(doc) + + async def _write_item(self, key: str, item: StoreItem) -> None: + + if key == "": + raise ValueError("CosmosDBStorage: Key cannot be empty.") + + doc = { + "id": sanitize_key( + key, self._config.key_suffix, self._config.compatibility_mode + ), + "realId": key, + "document": item.store_item_to_json(), + } + await self._container.upsert_item(body=doc) + + async def _delete_item(self, key: str) -> None: + + if key == "": + raise ValueError("CosmosDBStorage: Key cannot be empty.") + + escaped_key: str = sanitize_key( + key, self._config.key_suffix, self._config.compatibility_mode + ) + await ignore_error(self._container.delete_item( + escaped_key, self._get_partition_key(escaped_key) + ), is_status_code_error(404)) + + async def initialize(self) -> None: + if not self._container: + with self._lock: + # in case another thread attempted to initialize just before acquiring the lock + if self._container: return + + if not self._database: + self._database = await self._client.create_database_if_not_exists( + self._config.database_id + ) + + partition_key = { + "paths": ["/id"], + "kind": documents.PartitionKind.Hash, + } + try: + self._container = await self._database.create_container( + self._config.container_id, + partition_key, + offer_throughput=self._config.container_throughput, + ) + except cosmos_exceptions.CosmosHttpResponseError as err: + if err.status_code == http_constants.StatusCodes.CONFLICT: + self._container = self._database.get_container_client( + self._config.container_id + ) + properties = await self._container.read() + if "partitionKey" not in properties: + self._compatability_mode_partition_key = True + else: + paths = properties["partitionKey"]["paths"] + if "/partitionKey" in paths: + self._compatability_mode_partition_key = True + elif "/id" not in paths: + raise Exception( + f"Custom Partition Key Paths are not supported. {self._config.container_id} " + "has a custom Partition Key Path of {paths[0]}." + ) + else: + raise err + + def _get_partition_key(self, key: str) -> str: + return "" if self._compatability_mode_partition_key else key diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py new file mode 100644 index 00000000..2042887b --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py @@ -0,0 +1,88 @@ +import json +from typing import Union + +from azure.core.credentials import TokenCredential + +from .key_ops import sanitize_key + +class CosmosDBStorageConfig: + """The class for partitioned CosmosDB configuration for the Azure Bot Framework.""" + + def __init__( + self, + cosmos_db_endpoint: str = "", + auth_key: str = "", + database_id: str = "", + container_id: str = "", + cosmos_client_options: dict = None, + container_throughput: int = 0, + key_suffix: str = "", + compatibility_mode: bool = False, + url: str = "", + credential: Union[TokenCredential, None] = None, + **kwargs, + ): + """Create the Config object. + + :param cosmos_db_endpoint: The CosmosDB endpoint. + :param auth_key: The authentication key for Cosmos DB. + :param database_id: The database identifier for Cosmos DB instance. + :param container_id: The container identifier. + :param cosmos_client_options: The options for the CosmosClient. Currently only supports connection_policy and + consistency_level + :param container_throughput: The throughput set when creating the Container. Defaults to 400. + :param key_suffix: The suffix to be added to every key. The keySuffix must contain only valid ComosDb + key characters. (e.g. not: '\\', '?', '/', '#', '*') + :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb + max key length of 255. + :return CosmosDBConfig: + """ + config_file: str = kwargs.get("filename", "") + if config_file: + kwargs = json.load(open(config_file)) + self.cosmos_db_endpoint: str = cosmos_db_endpoint or kwargs.get( + "cosmos_db_endpoint", "" + ) + self.auth_key: str = auth_key or kwargs.get("auth_key", "") + self.database_id: str = database_id or kwargs.get("database_id", "") + self.container_id: str = container_id or kwargs.get("container_id", "") + self.cosmos_client_options: dict = cosmos_client_options or kwargs.get( + "cosmos_client_options", {} + ) + self.container_throughput: int = container_throughput or kwargs.get( + "container_throughput", 400 + ) + self.key_suffix: str = key_suffix or kwargs.get("key_suffix", "") + self.compatibility_mode: bool = compatibility_mode or kwargs.get( + "compatibility_mode", False + ) + self.url = url or kwargs.get("url", "") + self.credential: Union[TokenCredential, None] = credential + + @staticmethod + def validate_cosmos_db_config(config: "CosmosDBStorageConfig") -> None: + if not config: + raise ValueError("CosmosDBStorage: CosmosDBConfig is required.") + if not config.cosmos_db_endpoint: + raise ValueError("CosmosDBStorage: cosmos_db_endpoint is required.") + if not config.auth_key: + raise ValueError("CosmosDBStorage: auth_key is required.") + if not config.database_id: + raise ValueError("CosmosDBStorage: database_id is required.") + if not config.container_id: + raise ValueError("CosmosDBStorage: container_id is required.") + + CosmosDBStorageConfig._validate_suffix(config) + + @staticmethod + def _validate_suffix(config: "CosmosDBStorageConfig") -> None: + if config.key_suffix: + if config.compatibility_mode: + raise ValueError( + "compatibilityMode cannot be true while using a keySuffix." + ) + suffix_escaped: str = sanitize_key(config.key_suffix) + if suffix_escaped != config.key_suffix: + raise ValueError( + f"Cannot use invalid Row Key characters: {config.key_suffix} in keySuffix." + ) \ No newline at end of file diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py new file mode 100644 index 00000000..e2b0a8db --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py @@ -0,0 +1,37 @@ +from hashlib import sha256 + +def sanitize_key( + key: str, key_suffix: str = "", compatibility_mode: bool = True +) -> str: + """Return the sanitized key. + + Replace characters that are not allowed in keys in Cosmos. + + :param key: The provided key to be escaped. + :param key_suffix: The string to add a the end of all RowKeys. + :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb + max key length of 255. This behavior can be overridden by setting + cosmosdb_config.compatibility_mode to False. + :return str: + """ + # forbidden characters + bad_chars: list[str] = ["\\", "?", "/", "#", "\t", "\n", "\r", "*"] + + # replace those with with '*' and the + # Unicode code point of the character and return the new string + key = "".join(map(lambda x: "*" + str(ord(x)) if x in bad_chars else x, key)) + return truncate_key(f"{key}{key_suffix}", compatibility_mode) + +def truncate_key(key: str, compatibility_mode: bool = True) -> str: + max_key_len: int = 255 + + if not compatibility_mode: + return key + + if len(key) > max_key_len: + aux_hash = sha256(key.encode("utf-8")) + aux_hex = aux_hash.hexdigest() + + key = key[0 : max_key_len - len(aux_hex)] + aux_hex + + return key \ No newline at end of file diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/test.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/test.py new file mode 100644 index 00000000..afd2b476 --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/test.py @@ -0,0 +1,419 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import unittest +from unittest.mock import Mock, AsyncMock, patch, mock_open +import json +from typing import Dict, Any + +from azure.cosmos import documents, http_constants +from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy +import azure.cosmos.exceptions as cosmos_exceptions + +from microsoft.agents.storage import StoreItem +from cosmos import CosmosDBConfig, CosmosDBStorage, sanitize_key, truncate_key + + +class MockStoreItem(StoreItem): + """Mock StoreItem for testing""" + + def __init__(self, data: Dict[str, Any] = None): + self.data = data or {} + self.e_tag = "*" + + @classmethod + def from_json_to_store_item(cls, json_data: Dict[str, Any]) -> 'MockStoreItem': + item = cls() + item.data = json_data + return item + + def store_item_to_json(self) -> Dict[str, Any]: + return self.data + + +class TestCosmosDBConfig(unittest.TestCase): + """Test cases for CosmosDBConfig class""" + + def test_config_initialization_with_parameters(self): + """Test basic configuration initialization""" + config = CosmosDBConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db", + container_id="test_container", + container_throughput=800, + key_suffix="_test", + compatibility_mode=True + ) + + self.assertEqual(config.cosmos_db_endpoint, "https://test.documents.azure.com:443/") + self.assertEqual(config.auth_key, "test_key") + self.assertEqual(config.database_id, "test_db") + self.assertEqual(config.container_id, "test_container") + self.assertEqual(config.container_throughput, 800) + self.assertEqual(config.key_suffix, "_test") + self.assertTrue(config.compatibility_mode) + + @patch("builtins.open", new_callable=mock_open, read_data='{"cosmos_db_endpoint": "https://file.test.com", "auth_key": "file_key"}') + def test_config_initialization_from_file(self, mock_file): + """Test configuration initialization from JSON file""" + config = CosmosDBConfig(filename="test_config.json") + + self.assertEqual(config.cosmos_db_endpoint, "https://file.test.com") + self.assertEqual(config.auth_key, "file_key") + mock_file.assert_called_once_with("test_config.json") + + def test_config_parameter_override_file(self): + """Test that direct parameters override file parameters""" + with patch("builtins.open", new_callable=mock_open, read_data='{"cosmos_db_endpoint": "https://file.test.com"}'): + config = CosmosDBConfig( + cosmos_db_endpoint="https://override.test.com", + filename="test_config.json" + ) + + self.assertEqual(config.cosmos_db_endpoint, "https://override.test.com") + + +class TestKeySanitization(unittest.TestCase): + """Test cases for key sanitization functions""" + + def test_sanitize_key_normal_string(self): + """Test sanitization of normal string""" + result = sanitize_key("normal_key", compatibility_mode=False) + self.assertEqual(result, "normal_key") + + def test_sanitize_key_with_forbidden_chars(self): + """Test sanitization of string with forbidden characters""" + result = sanitize_key("key\\with?bad/chars#", compatibility_mode=False) + expected = "key*92with*63bad*47chars*35" + self.assertEqual(result, expected) + + def test_sanitize_key_with_suffix(self): + """Test sanitization with key suffix""" + result = sanitize_key("test_key", key_suffix="_suffix", compatibility_mode=False) + self.assertEqual(result, "test_key_suffix") + + def test_sanitize_key_with_tabs_and_newlines(self): + """Test sanitization of tabs and newlines""" + result = sanitize_key("key\twith\nnewlines\rand\ttabs", compatibility_mode=False) + expected = "key*9with*10newlines*13and*9tabs" + self.assertEqual(result, expected) + + def test_truncate_key_short_string(self): + """Test truncation of short string""" + result = truncate_key("short_key", compatibility_mode=True) + self.assertEqual(result, "short_key") + + def test_truncate_key_long_string(self): + """Test truncation of long string""" + long_key = "a" * 300 + result = truncate_key(long_key, compatibility_mode=True) + + # Should be exactly 255 characters + self.assertEqual(len(result), 255) + # Should end with a hash + self.assertTrue(result.endswith("aa" * 16)) # SHA256 hex is 64 chars, but truncated + + def test_truncate_key_disabled(self): + """Test truncation when compatibility mode is disabled""" + long_key = "a" * 300 + result = truncate_key(long_key, compatibility_mode=False) + self.assertEqual(result, long_key) + self.assertEqual(len(result), 300) + + +class TestCosmosDBStorage(unittest.TestCase): + """Test cases for CosmosDBStorage class""" + + def setUp(self): + """Set up test fixtures""" + self.config = CosmosDBConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db", + container_id="test_container" + ) + + self.mock_client = Mock(spec=CosmosClient) + self.mock_database = Mock(spec=DatabaseProxy) + self.mock_container = Mock(spec=ContainerProxy) + + def test_config_validation_success(self): + """Test successful configuration validation""" + # Should not raise any exception + CosmosDBStorage._validate_config(self.config) + + def test_config_validation_missing_config(self): + """Test validation with missing config""" + with self.assertRaises(ValueError) as cm: + CosmosDBStorage._validate_config(None) + self.assertIn("CosmosDBConfig is required", str(cm.exception)) + + def test_config_validation_missing_endpoint(self): + """Test validation with missing endpoint""" + config = CosmosDBConfig(auth_key="key", database_id="db", container_id="container") + with self.assertRaises(ValueError) as cm: + CosmosDBStorage._validate_config(config) + self.assertIn("cosmos_db_endpoint is required", str(cm.exception)) + + def test_config_validation_missing_auth_key(self): + """Test validation with missing auth key""" + config = CosmosDBConfig(cosmos_db_endpoint="endpoint", database_id="db", container_id="container") + with self.assertRaises(ValueError) as cm: + CosmosDBStorage._validate_config(config) + self.assertIn("auth_key is required", str(cm.exception)) + + def test_suffix_validation_success(self): + """Test successful suffix validation""" + config = CosmosDBConfig(key_suffix="valid_suffix", compatibility_mode=False) + # Should not raise any exception + CosmosDBStorage._validate_suffix(config) + + def test_suffix_validation_with_compatibility_mode(self): + """Test suffix validation fails with compatibility mode""" + config = CosmosDBConfig(key_suffix="suffix", compatibility_mode=True) + with self.assertRaises(ValueError) as cm: + CosmosDBStorage._validate_suffix(config) + self.assertIn("compatibilityMode cannot be true while using a keySuffix", str(cm.exception)) + + def test_suffix_validation_invalid_characters(self): + """Test suffix validation with invalid characters""" + config = CosmosDBConfig(key_suffix="bad\\suffix", compatibility_mode=False) + with self.assertRaises(ValueError) as cm: + CosmosDBStorage._validate_suffix(config) + self.assertIn("Cannot use invalid Row Key characters", str(cm.exception)) + + @patch('cosmos.CosmosClient') + def test_create_client_success(self, mock_cosmos_client): + """Test successful client creation""" + storage = CosmosDBStorage(self.config) + client = storage._create_client() + + mock_cosmos_client.assert_called_once() + args, kwargs = mock_cosmos_client.call_args + self.assertEqual(args[0], self.config.cosmos_db_endpoint) + self.assertEqual(args[1], self.config.auth_key) + + @patch('cosmos.CosmosClient') + def test_storage_initialization(self, mock_cosmos_client): + """Test storage initialization""" + mock_cosmos_client.return_value = self.mock_client + + storage = CosmosDBStorage(self.config) + + self.assertEqual(storage._config, self.config) + self.assertIsNotNone(storage._client) + self.assertIsNone(storage._database) + self.assertIsNone(storage._container) + self.assertFalse(storage._compatability_mode_partition_key) + + +class TestCosmosDBStorageAsync(unittest.IsolatedAsyncioTestCase): + """Async test cases for CosmosDBStorage class""" + + async def asyncSetUp(self): + """Set up async test fixtures""" + self.config = CosmosDBConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db", + container_id="test_container" + ) + + self.mock_client = Mock(spec=CosmosClient) + self.mock_database = Mock(spec=DatabaseProxy) + self.mock_container = Mock(spec=ContainerProxy) + + # Configure async mocks + self.mock_container.read_item = AsyncMock() + self.mock_container.upsert_item = AsyncMock() + self.mock_container.delete_item = AsyncMock() + self.mock_container.read = Mock() + + self.mock_database.create_container = Mock() + self.mock_database.get_container_client = Mock() + + self.mock_client.create_database_if_not_exists = Mock() + + @patch('cosmos.CosmosClient') + async def test_read_success(self, mock_cosmos_client): + """Test successful read operation""" + mock_cosmos_client.return_value = self.mock_client + self.mock_client.create_database_if_not_exists.return_value = self.mock_database + self.mock_database.create_container.return_value = self.mock_container + + # Mock the read response + mock_response = { + "realId": "test_key", + "document": {"data": "test_data"} + } + self.mock_container.read_item.return_value = mock_response + + storage = CosmosDBStorage(self.config) + result = await storage.read(["test_key"], target_cls=MockStoreItem) + + self.assertIn("test_key", result) + self.assertEqual(result["test_key"].data, {"data": "test_data"}) + + @patch('cosmos.CosmosClient') + async def test_read_empty_keys(self, mock_cosmos_client): + """Test read with empty keys list""" + mock_cosmos_client.return_value = self.mock_client + + storage = CosmosDBStorage(self.config) + + with self.assertRaises(ValueError) as cm: + await storage.read([]) + self.assertIn("keys cannot be None or empty", str(cm.exception)) + + @patch('cosmos.CosmosClient') + async def test_write_success(self, mock_cosmos_client): + """Test successful write operation""" + mock_cosmos_client.return_value = self.mock_client + self.mock_client.create_database_if_not_exists.return_value = self.mock_database + self.mock_database.create_container.return_value = self.mock_container + + storage = CosmosDBStorage(self.config) + + test_item = MockStoreItem({"test": "data"}) + changes = {"test_key": test_item} + + await storage.write(changes) + + # Verify upsert_item was called + self.mock_container.upsert_item.assert_called_once() + args, kwargs = self.mock_container.upsert_item.call_args + + # Check the document structure + doc = kwargs.get('body') or args[0] if args else None + self.assertIsNotNone(doc) + self.assertEqual(doc["realId"], "test_key") + self.assertEqual(doc["document"], {"test": "data"}) + + @patch('cosmos.CosmosClient') + async def test_write_empty_changes(self, mock_cosmos_client): + """Test write with empty changes""" + mock_cosmos_client.return_value = self.mock_client + + storage = CosmosDBStorage(self.config) + + with self.assertRaises(Exception) as cm: + await storage.write({}) + self.assertIn("Changes are required when writing", str(cm.exception)) + + @patch('cosmos.CosmosClient') + async def test_delete_success(self, mock_cosmos_client): + """Test successful delete operation""" + mock_cosmos_client.return_value = self.mock_client + self.mock_client.create_database_if_not_exists.return_value = self.mock_database + self.mock_database.create_container.return_value = self.mock_container + + storage = CosmosDBStorage(self.config) + + await storage.delete(["test_key1", "test_key2"]) + + # Verify delete_item was called for each key + self.assertEqual(self.mock_container.delete_item.call_count, 2) + + @patch('cosmos.CosmosClient') + async def test_delete_none_keys(self, mock_cosmos_client): + """Test delete with None keys""" + mock_cosmos_client.return_value = self.mock_client + + storage = CosmosDBStorage(self.config) + + with self.assertRaises(ValueError) as cm: + await storage.delete(None) + self.assertIn("keys parameter can't be null", str(cm.exception)) + + @patch('cosmos.CosmosClient') + async def test_initialize_new_container(self, mock_cosmos_client): + """Test initialization with new container creation""" + mock_cosmos_client.return_value = self.mock_client + self.mock_client.create_database_if_not_exists.return_value = self.mock_database + self.mock_database.create_container.return_value = self.mock_container + + storage = CosmosDBStorage(self.config) + await storage.initialize() + + # Verify database and container creation + self.mock_client.create_database_if_not_exists.assert_called_once_with("test_db") + self.mock_database.create_container.assert_called_once() + + # Check partition key configuration + args, kwargs = self.mock_database.create_container.call_args + partition_key = args[1] + self.assertEqual(partition_key["paths"], ["/id"]) + self.assertEqual(partition_key["kind"], documents.PartitionKind.Hash) + + @patch('cosmos.CosmosClient') + async def test_initialize_existing_container_conflict(self, mock_cosmos_client): + """Test initialization with existing container (conflict scenario)""" + mock_cosmos_client.return_value = self.mock_client + self.mock_client.create_database_if_not_exists.return_value = self.mock_database + + # Mock conflict exception + conflict_error = cosmos_exceptions.CosmosHttpResponseError( + status_code=http_constants.StatusCodes.CONFLICT + ) + self.mock_database.create_container.side_effect = conflict_error + self.mock_database.get_container_client.return_value = self.mock_container + + # Mock container properties + self.mock_container.read.return_value = { + "partitionKey": {"paths": ["/id"]} + } + + storage = CosmosDBStorage(self.config) + await storage.initialize() + + # Verify fallback to get existing container + self.mock_database.get_container_client.assert_called_once_with("test_container") + self.assertFalse(storage._compatability_mode_partition_key) + + @patch('cosmos.CosmosClient') + async def test_initialize_compatibility_mode_partition_key(self, mock_cosmos_client): + """Test initialization with compatibility mode partition key""" + mock_cosmos_client.return_value = self.mock_client + self.mock_client.create_database_if_not_exists.return_value = self.mock_database + + # Mock conflict exception + conflict_error = cosmos_exceptions.CosmosHttpResponseError( + status_code=http_constants.StatusCodes.CONFLICT + ) + self.mock_database.create_container.side_effect = conflict_error + self.mock_database.get_container_client.return_value = self.mock_container + + # Mock container properties with old partition key + self.mock_container.read.return_value = { + "partitionKey": {"paths": ["/partitionKey"]} + } + + storage = CosmosDBStorage(self.config) + await storage.initialize() + + # Should enable compatibility mode + self.assertTrue(storage._compatability_mode_partition_key) + + def test_get_partition_key_normal_mode(self): + """Test partition key generation in normal mode""" + with patch('cosmos.CosmosClient'): + storage = CosmosDBStorage(self.config) + storage._compatability_mode_partition_key = False + + result = storage._get_partition_key("test_key") + self.assertEqual(result, "test_key") + + def test_get_partition_key_compatibility_mode(self): + """Test partition key generation in compatibility mode""" + with patch('cosmos.CosmosClient'): + storage = CosmosDBStorage(self.config) + storage._compatability_mode_partition_key = True + + result = storage._get_partition_key("test_key") + self.assertEqual(result, "") + + +if __name__ == '__main__': + # Run the tests + unittest.main() \ No newline at end of file From 04b506a03a158fcaafae3238a80e03f46f7df570 Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Fri, 18 Jul 2025 09:21:35 -0700 Subject: [PATCH 02/13] Removed unnecessary file --- .../microsoft/agents/cosmos/test.py | 419 ------------------ 1 file changed, 419 deletions(-) delete mode 100644 libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/test.py diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/test.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/test.py deleted file mode 100644 index afd2b476..00000000 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/test.py +++ /dev/null @@ -1,419 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. - -import unittest -from unittest.mock import Mock, AsyncMock, patch, mock_open -import json -from typing import Dict, Any - -from azure.cosmos import documents, http_constants -from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy -import azure.cosmos.exceptions as cosmos_exceptions - -from microsoft.agents.storage import StoreItem -from cosmos import CosmosDBConfig, CosmosDBStorage, sanitize_key, truncate_key - - -class MockStoreItem(StoreItem): - """Mock StoreItem for testing""" - - def __init__(self, data: Dict[str, Any] = None): - self.data = data or {} - self.e_tag = "*" - - @classmethod - def from_json_to_store_item(cls, json_data: Dict[str, Any]) -> 'MockStoreItem': - item = cls() - item.data = json_data - return item - - def store_item_to_json(self) -> Dict[str, Any]: - return self.data - - -class TestCosmosDBConfig(unittest.TestCase): - """Test cases for CosmosDBConfig class""" - - def test_config_initialization_with_parameters(self): - """Test basic configuration initialization""" - config = CosmosDBConfig( - cosmos_db_endpoint="https://test.documents.azure.com:443/", - auth_key="test_key", - database_id="test_db", - container_id="test_container", - container_throughput=800, - key_suffix="_test", - compatibility_mode=True - ) - - self.assertEqual(config.cosmos_db_endpoint, "https://test.documents.azure.com:443/") - self.assertEqual(config.auth_key, "test_key") - self.assertEqual(config.database_id, "test_db") - self.assertEqual(config.container_id, "test_container") - self.assertEqual(config.container_throughput, 800) - self.assertEqual(config.key_suffix, "_test") - self.assertTrue(config.compatibility_mode) - - @patch("builtins.open", new_callable=mock_open, read_data='{"cosmos_db_endpoint": "https://file.test.com", "auth_key": "file_key"}') - def test_config_initialization_from_file(self, mock_file): - """Test configuration initialization from JSON file""" - config = CosmosDBConfig(filename="test_config.json") - - self.assertEqual(config.cosmos_db_endpoint, "https://file.test.com") - self.assertEqual(config.auth_key, "file_key") - mock_file.assert_called_once_with("test_config.json") - - def test_config_parameter_override_file(self): - """Test that direct parameters override file parameters""" - with patch("builtins.open", new_callable=mock_open, read_data='{"cosmos_db_endpoint": "https://file.test.com"}'): - config = CosmosDBConfig( - cosmos_db_endpoint="https://override.test.com", - filename="test_config.json" - ) - - self.assertEqual(config.cosmos_db_endpoint, "https://override.test.com") - - -class TestKeySanitization(unittest.TestCase): - """Test cases for key sanitization functions""" - - def test_sanitize_key_normal_string(self): - """Test sanitization of normal string""" - result = sanitize_key("normal_key", compatibility_mode=False) - self.assertEqual(result, "normal_key") - - def test_sanitize_key_with_forbidden_chars(self): - """Test sanitization of string with forbidden characters""" - result = sanitize_key("key\\with?bad/chars#", compatibility_mode=False) - expected = "key*92with*63bad*47chars*35" - self.assertEqual(result, expected) - - def test_sanitize_key_with_suffix(self): - """Test sanitization with key suffix""" - result = sanitize_key("test_key", key_suffix="_suffix", compatibility_mode=False) - self.assertEqual(result, "test_key_suffix") - - def test_sanitize_key_with_tabs_and_newlines(self): - """Test sanitization of tabs and newlines""" - result = sanitize_key("key\twith\nnewlines\rand\ttabs", compatibility_mode=False) - expected = "key*9with*10newlines*13and*9tabs" - self.assertEqual(result, expected) - - def test_truncate_key_short_string(self): - """Test truncation of short string""" - result = truncate_key("short_key", compatibility_mode=True) - self.assertEqual(result, "short_key") - - def test_truncate_key_long_string(self): - """Test truncation of long string""" - long_key = "a" * 300 - result = truncate_key(long_key, compatibility_mode=True) - - # Should be exactly 255 characters - self.assertEqual(len(result), 255) - # Should end with a hash - self.assertTrue(result.endswith("aa" * 16)) # SHA256 hex is 64 chars, but truncated - - def test_truncate_key_disabled(self): - """Test truncation when compatibility mode is disabled""" - long_key = "a" * 300 - result = truncate_key(long_key, compatibility_mode=False) - self.assertEqual(result, long_key) - self.assertEqual(len(result), 300) - - -class TestCosmosDBStorage(unittest.TestCase): - """Test cases for CosmosDBStorage class""" - - def setUp(self): - """Set up test fixtures""" - self.config = CosmosDBConfig( - cosmos_db_endpoint="https://test.documents.azure.com:443/", - auth_key="test_key", - database_id="test_db", - container_id="test_container" - ) - - self.mock_client = Mock(spec=CosmosClient) - self.mock_database = Mock(spec=DatabaseProxy) - self.mock_container = Mock(spec=ContainerProxy) - - def test_config_validation_success(self): - """Test successful configuration validation""" - # Should not raise any exception - CosmosDBStorage._validate_config(self.config) - - def test_config_validation_missing_config(self): - """Test validation with missing config""" - with self.assertRaises(ValueError) as cm: - CosmosDBStorage._validate_config(None) - self.assertIn("CosmosDBConfig is required", str(cm.exception)) - - def test_config_validation_missing_endpoint(self): - """Test validation with missing endpoint""" - config = CosmosDBConfig(auth_key="key", database_id="db", container_id="container") - with self.assertRaises(ValueError) as cm: - CosmosDBStorage._validate_config(config) - self.assertIn("cosmos_db_endpoint is required", str(cm.exception)) - - def test_config_validation_missing_auth_key(self): - """Test validation with missing auth key""" - config = CosmosDBConfig(cosmos_db_endpoint="endpoint", database_id="db", container_id="container") - with self.assertRaises(ValueError) as cm: - CosmosDBStorage._validate_config(config) - self.assertIn("auth_key is required", str(cm.exception)) - - def test_suffix_validation_success(self): - """Test successful suffix validation""" - config = CosmosDBConfig(key_suffix="valid_suffix", compatibility_mode=False) - # Should not raise any exception - CosmosDBStorage._validate_suffix(config) - - def test_suffix_validation_with_compatibility_mode(self): - """Test suffix validation fails with compatibility mode""" - config = CosmosDBConfig(key_suffix="suffix", compatibility_mode=True) - with self.assertRaises(ValueError) as cm: - CosmosDBStorage._validate_suffix(config) - self.assertIn("compatibilityMode cannot be true while using a keySuffix", str(cm.exception)) - - def test_suffix_validation_invalid_characters(self): - """Test suffix validation with invalid characters""" - config = CosmosDBConfig(key_suffix="bad\\suffix", compatibility_mode=False) - with self.assertRaises(ValueError) as cm: - CosmosDBStorage._validate_suffix(config) - self.assertIn("Cannot use invalid Row Key characters", str(cm.exception)) - - @patch('cosmos.CosmosClient') - def test_create_client_success(self, mock_cosmos_client): - """Test successful client creation""" - storage = CosmosDBStorage(self.config) - client = storage._create_client() - - mock_cosmos_client.assert_called_once() - args, kwargs = mock_cosmos_client.call_args - self.assertEqual(args[0], self.config.cosmos_db_endpoint) - self.assertEqual(args[1], self.config.auth_key) - - @patch('cosmos.CosmosClient') - def test_storage_initialization(self, mock_cosmos_client): - """Test storage initialization""" - mock_cosmos_client.return_value = self.mock_client - - storage = CosmosDBStorage(self.config) - - self.assertEqual(storage._config, self.config) - self.assertIsNotNone(storage._client) - self.assertIsNone(storage._database) - self.assertIsNone(storage._container) - self.assertFalse(storage._compatability_mode_partition_key) - - -class TestCosmosDBStorageAsync(unittest.IsolatedAsyncioTestCase): - """Async test cases for CosmosDBStorage class""" - - async def asyncSetUp(self): - """Set up async test fixtures""" - self.config = CosmosDBConfig( - cosmos_db_endpoint="https://test.documents.azure.com:443/", - auth_key="test_key", - database_id="test_db", - container_id="test_container" - ) - - self.mock_client = Mock(spec=CosmosClient) - self.mock_database = Mock(spec=DatabaseProxy) - self.mock_container = Mock(spec=ContainerProxy) - - # Configure async mocks - self.mock_container.read_item = AsyncMock() - self.mock_container.upsert_item = AsyncMock() - self.mock_container.delete_item = AsyncMock() - self.mock_container.read = Mock() - - self.mock_database.create_container = Mock() - self.mock_database.get_container_client = Mock() - - self.mock_client.create_database_if_not_exists = Mock() - - @patch('cosmos.CosmosClient') - async def test_read_success(self, mock_cosmos_client): - """Test successful read operation""" - mock_cosmos_client.return_value = self.mock_client - self.mock_client.create_database_if_not_exists.return_value = self.mock_database - self.mock_database.create_container.return_value = self.mock_container - - # Mock the read response - mock_response = { - "realId": "test_key", - "document": {"data": "test_data"} - } - self.mock_container.read_item.return_value = mock_response - - storage = CosmosDBStorage(self.config) - result = await storage.read(["test_key"], target_cls=MockStoreItem) - - self.assertIn("test_key", result) - self.assertEqual(result["test_key"].data, {"data": "test_data"}) - - @patch('cosmos.CosmosClient') - async def test_read_empty_keys(self, mock_cosmos_client): - """Test read with empty keys list""" - mock_cosmos_client.return_value = self.mock_client - - storage = CosmosDBStorage(self.config) - - with self.assertRaises(ValueError) as cm: - await storage.read([]) - self.assertIn("keys cannot be None or empty", str(cm.exception)) - - @patch('cosmos.CosmosClient') - async def test_write_success(self, mock_cosmos_client): - """Test successful write operation""" - mock_cosmos_client.return_value = self.mock_client - self.mock_client.create_database_if_not_exists.return_value = self.mock_database - self.mock_database.create_container.return_value = self.mock_container - - storage = CosmosDBStorage(self.config) - - test_item = MockStoreItem({"test": "data"}) - changes = {"test_key": test_item} - - await storage.write(changes) - - # Verify upsert_item was called - self.mock_container.upsert_item.assert_called_once() - args, kwargs = self.mock_container.upsert_item.call_args - - # Check the document structure - doc = kwargs.get('body') or args[0] if args else None - self.assertIsNotNone(doc) - self.assertEqual(doc["realId"], "test_key") - self.assertEqual(doc["document"], {"test": "data"}) - - @patch('cosmos.CosmosClient') - async def test_write_empty_changes(self, mock_cosmos_client): - """Test write with empty changes""" - mock_cosmos_client.return_value = self.mock_client - - storage = CosmosDBStorage(self.config) - - with self.assertRaises(Exception) as cm: - await storage.write({}) - self.assertIn("Changes are required when writing", str(cm.exception)) - - @patch('cosmos.CosmosClient') - async def test_delete_success(self, mock_cosmos_client): - """Test successful delete operation""" - mock_cosmos_client.return_value = self.mock_client - self.mock_client.create_database_if_not_exists.return_value = self.mock_database - self.mock_database.create_container.return_value = self.mock_container - - storage = CosmosDBStorage(self.config) - - await storage.delete(["test_key1", "test_key2"]) - - # Verify delete_item was called for each key - self.assertEqual(self.mock_container.delete_item.call_count, 2) - - @patch('cosmos.CosmosClient') - async def test_delete_none_keys(self, mock_cosmos_client): - """Test delete with None keys""" - mock_cosmos_client.return_value = self.mock_client - - storage = CosmosDBStorage(self.config) - - with self.assertRaises(ValueError) as cm: - await storage.delete(None) - self.assertIn("keys parameter can't be null", str(cm.exception)) - - @patch('cosmos.CosmosClient') - async def test_initialize_new_container(self, mock_cosmos_client): - """Test initialization with new container creation""" - mock_cosmos_client.return_value = self.mock_client - self.mock_client.create_database_if_not_exists.return_value = self.mock_database - self.mock_database.create_container.return_value = self.mock_container - - storage = CosmosDBStorage(self.config) - await storage.initialize() - - # Verify database and container creation - self.mock_client.create_database_if_not_exists.assert_called_once_with("test_db") - self.mock_database.create_container.assert_called_once() - - # Check partition key configuration - args, kwargs = self.mock_database.create_container.call_args - partition_key = args[1] - self.assertEqual(partition_key["paths"], ["/id"]) - self.assertEqual(partition_key["kind"], documents.PartitionKind.Hash) - - @patch('cosmos.CosmosClient') - async def test_initialize_existing_container_conflict(self, mock_cosmos_client): - """Test initialization with existing container (conflict scenario)""" - mock_cosmos_client.return_value = self.mock_client - self.mock_client.create_database_if_not_exists.return_value = self.mock_database - - # Mock conflict exception - conflict_error = cosmos_exceptions.CosmosHttpResponseError( - status_code=http_constants.StatusCodes.CONFLICT - ) - self.mock_database.create_container.side_effect = conflict_error - self.mock_database.get_container_client.return_value = self.mock_container - - # Mock container properties - self.mock_container.read.return_value = { - "partitionKey": {"paths": ["/id"]} - } - - storage = CosmosDBStorage(self.config) - await storage.initialize() - - # Verify fallback to get existing container - self.mock_database.get_container_client.assert_called_once_with("test_container") - self.assertFalse(storage._compatability_mode_partition_key) - - @patch('cosmos.CosmosClient') - async def test_initialize_compatibility_mode_partition_key(self, mock_cosmos_client): - """Test initialization with compatibility mode partition key""" - mock_cosmos_client.return_value = self.mock_client - self.mock_client.create_database_if_not_exists.return_value = self.mock_database - - # Mock conflict exception - conflict_error = cosmos_exceptions.CosmosHttpResponseError( - status_code=http_constants.StatusCodes.CONFLICT - ) - self.mock_database.create_container.side_effect = conflict_error - self.mock_database.get_container_client.return_value = self.mock_container - - # Mock container properties with old partition key - self.mock_container.read.return_value = { - "partitionKey": {"paths": ["/partitionKey"]} - } - - storage = CosmosDBStorage(self.config) - await storage.initialize() - - # Should enable compatibility mode - self.assertTrue(storage._compatability_mode_partition_key) - - def test_get_partition_key_normal_mode(self): - """Test partition key generation in normal mode""" - with patch('cosmos.CosmosClient'): - storage = CosmosDBStorage(self.config) - storage._compatability_mode_partition_key = False - - result = storage._get_partition_key("test_key") - self.assertEqual(result, "test_key") - - def test_get_partition_key_compatibility_mode(self): - """Test partition key generation in compatibility mode""" - with patch('cosmos.CosmosClient'): - storage = CosmosDBStorage(self.config) - storage._compatability_mode_partition_key = True - - result = storage._get_partition_key("test_key") - self.assertEqual(result, "") - - -if __name__ == '__main__': - # Run the tests - unittest.main() \ No newline at end of file From 6f9b88b7d311a30f8ec235cedcce650a9886ef2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Fri, 18 Jul 2025 10:12:35 -0700 Subject: [PATCH 03/13] Added quick storage testing functionality for debugging and improved documentation --- .../agents/storage/error_handling.py | 8 +++ .../microsoft/agents/storage/storage.py | 5 +- .../agents/storage/storage_test_utils.py | 61 ++++++++++++++++++- 3 files changed, 70 insertions(+), 4 deletions(-) diff --git a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/error_handling.py b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/error_handling.py index 396c8d2f..5bcd89b4 100644 --- a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/error_handling.py +++ b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/error_handling.py @@ -7,8 +7,13 @@ async def ignore_error(promise: Awaitable, ignore_error_filter: error_filter): """ Ignores errors based on the provided filter function. + promise: the awaitable to execute ignore_error_filter: a function that takes an Exception and returns True if the error should be + ignored, False otherwise. + + Returns the result of the promise if successful, or None if the error is ignored. + Raises the error if it is not ignored. """ try: return await promise @@ -21,6 +26,9 @@ async def ignore_error(promise: Awaitable, ignore_error_filter: error_filter): def is_status_code_error(*ignored_codes: list[int]) -> error_filter: """ Creates an error filter function that ignores errors with specific status codes. + + ignored_codes: a list of status codes to ignore + Returns a function that takes an Exception and returns True if the error's status code is in ignored_codes. """ def func(err: Exception) -> bool: diff --git a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage.py b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage.py index e9b0fcdd..4a71d939 100644 --- a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage.py +++ b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage.py @@ -40,7 +40,10 @@ async def delete(self, keys: list[str]) -> None: class AsyncStorageBase(Storage): - """Base class for asynchronous storage implementations.""" + """Base class for asynchronous storage implementations with operations + that work on single items. The bulk operations are implemented in terms + of the single-item operations. + """ async def initialize(self) -> None: """Initializes the storage container""" diff --git a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py index 9f805dae..8c314392 100644 --- a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py +++ b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py @@ -133,10 +133,14 @@ def delete(self, keys: list[str]) -> None: async def equals(self, other) -> bool: """ - Compare the items for all keys seenby this mock instance. + Compare the items for all keys seen by this mock instance. + + Note: This is an extra safety measure, and I've made the executive decision to not test this method itself - as it is not the main focus of the test suite. + because passing tests with calls to this method + is also dependent on the correctness of other + aspects, based on the other assertions in the tests. """ for key in self._key_history: if key not in self._memory: @@ -155,6 +159,7 @@ async def equals(self, other) -> bool: class StorageTestsCommon(ABC): + """Common fixtures for Storage implementations.""" KEY_LIST = [ "f", @@ -211,8 +216,20 @@ def changes(self, request): class CRUDStorageTests(StorageTestsCommon): + """Tests for Storage implementations that support CRUD operations. + + To use, subclass and implement the `storage` method. + """ - async def storage(self, initial_data=None, existing=False): + async def storage(self, initial_data=None, existing=False) -> StorageMock: + """Return a StorageMock instance to be tested. + :param initial_data: The initial data to populate the storage with. + :param existing: If True, the storage instance should connect to an existing store. + + When testing your own storage implementation, ensure that you also extend + StorageMock with a wrapper around your storage implementation. You would then + return an instance of that wrapper here. + """ raise NotImplementedError("Subclasses must implement this") @pytest.mark.asyncio @@ -452,3 +469,41 @@ async def test_flow(self): gc.collect() storage_alt = await self.storage(existing=True) assert await baseline_storage.equals(storage_alt) + +class QuickCRUDStorageTests(CRUDStorageTests): + """Reduced set of permutations for quicker tests. Useful for debugging.""" + + KEY_LIST = ([ + "\\?/#\t\n\r*", "test.txt" + ]) + + READ_KEY_LIST = KEY_LIST + ["nonexistent_key"] + + STATE_LIST = [ + { key: MockStoreItem({"id": key, "value": f"value{key}"}) for key in KEY_LIST } + ] + + @pytest.fixture(params=STATE_LIST) + def initial_state(self, request): + return request.param + + @pytest.fixture(params=KEY_LIST) + def key(self, request): + return request.param + + @pytest.fixture(params=[KEY_LIST]) + def keys(self, request): + return request.param + + @pytest.fixture(params=subsets(KEY_LIST, 2)) + def changes(self, request): + changes_obj = {} + keys = request.param + changes_obj["new_key"] = MockStoreItemB({"field": "new_value_for_new_key"}, True) + for i, key in enumerate(keys): + if i % 2 == 0: + changes_obj[key] = MockStoreItemB({"data": f"value{key}"}, (i // 2) % 2 == 0) + else: + changes_obj[key] = MockStoreItem({"id": key, "value": f"new_value_for_{key}"}) + changes_obj["new_key_2"] = MockStoreItem({"field": "new_value_for_new_key_2"}) + return changes_obj From fe9dae55db3c4c786691a5103ecf9cbf9453d430 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Fri, 18 Jul 2025 14:59:48 -0700 Subject: [PATCH 04/13] Simplified storage test setup and added unit test --- .../tests/test_blob_storage.py | 25 ++++++------ .../agents/storage/storage_test_utils.py | 38 +++++++------------ .../tests/test_memory_storage.py | 20 ++-------- 3 files changed, 29 insertions(+), 54 deletions(-) diff --git a/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py b/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py index 8bea6fce..cba9f88d 100644 --- a/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py +++ b/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py @@ -12,7 +12,6 @@ from microsoft.agents.storage.storage_test_utils import ( CRUDStorageTests, - StorageMock, StorageBaseline, MockStoreItem, MockStoreItemB, @@ -20,7 +19,6 @@ EMULATOR_RUNNING = False - async def blob_storage_instance(existing=False): # Default Azure Storage Emulator connection string @@ -56,7 +54,6 @@ async def blob_storage_instance(existing=False): storage = BlobStorage(blob_storage_config) return storage, container_client - @pytest_asyncio.fixture async def blob_storage(): @@ -68,16 +65,6 @@ async def blob_storage(): # teardown await container_client.delete_container() - -class BlobStorageMock(StorageMock): - - def __init__(self, blob_storage): - self.storage = blob_storage - - def get_backing_store(self): - return self.storage - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") class TestBlobStorage(CRUDStorageTests): @@ -90,7 +77,7 @@ async def storage(self, initial_data=None, existing=False): value_rep = json.dumps(value.store_item_to_json()) await container_client.upload_blob(name=key, data=value_rep, overwrite=True) - return BlobStorageMock(storage) + return storage @pytest.mark.asyncio async def test_initialize(self, blob_storage): @@ -104,6 +91,16 @@ async def test_initialize(self, blob_storage): "key": MockStoreItem({"id": "item", "value": "data"}) } + @pytest.mark.asyncio + async def test_external_change_is_visible(self): + blob_storage, container_client = await blob_storage_instance() + assert (await blob_storage.read(["key"], target_cls=MockStoreItem)) == MockStoreItem({"id": "item", "value": "data"}) + assert (await blob_storage.read(["key2"], target_cls=MockStoreItem)) == MockStoreItem({"id": "another_item", "value": "another_value"}) + await container_client.upload_blob(name="key", data=json.dumps({"id": "item", "value": "data"}), overwrite=True) + await container_client.upload_blob(name="key2", data=json.dumps({"id": "another_item", "value": "new_val"}), overwrite=True) + assert (await blob_storage.read(["key"], target_cls=MockStoreItem)) == MockStoreItem({"id": "item", "value": "data"}) + assert (await blob_storage.read(["key2"], target_cls=MockStoreItem)) == MockStoreItem({"id": "another_item", "value": "new_val"}) + @pytest.mark.asyncio async def test_blob_storage_flow_existing_container_and_persistence(self): diff --git a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py index 8c314392..edc603dd 100644 --- a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py +++ b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py @@ -9,7 +9,6 @@ from ._type_aliases import JSON from .memory_storage import MemoryStorage - class MockStoreItem(StoreItem): """Test implementation of StoreItem for testing purposes""" @@ -92,22 +91,6 @@ def subsets(lst, n=-1): return subsets -class StorageMock(ABC): - """A mock wrapper around a Storage implementation to be used in tests.""" - - def get_backing_store(self) -> Storage: - raise NotImplementedError("Subclasses must implement this") - - async def read(self, *args, **kwargs): - return await self.get_backing_store().read(*args, **kwargs) - - async def write(self, *args, **kwargs): - return await self.get_backing_store().write(*args, **kwargs) - - async def delete(self, *args, **kwargs): - return await self.get_backing_store().delete(*args, **kwargs) - - # bootstrapping class to compare against # if this class is correct, then the tests are correct class StorageBaseline(Storage): @@ -221,14 +204,10 @@ class CRUDStorageTests(StorageTestsCommon): To use, subclass and implement the `storage` method. """ - async def storage(self, initial_data=None, existing=False) -> StorageMock: - """Return a StorageMock instance to be tested. + async def storage(self, initial_data=None, existing=False) -> Storage: + """Return a Storage instance to be tested. :param initial_data: The initial data to populate the storage with. :param existing: If True, the storage instance should connect to an existing store. - - When testing your own storage implementation, ensure that you also extend - StorageMock with a wrapper around your storage implementation. You would then - return an instance of that wrapper here. """ raise NotImplementedError("Subclasses must implement this") @@ -463,7 +442,7 @@ async def test_flow(self): await storage.read(["key_b"], target_cls=MockStoreItemB) assert await baseline_storage.equals(storage) - if not isinstance(storage.get_backing_store(), MemoryStorage): + if not isinstance(storage, MemoryStorage): # if not memory storage, then items should persist del storage gc.collect() @@ -507,3 +486,14 @@ def changes(self, request): changes_obj[key] = MockStoreItem({"id": key, "value": f"new_value_for_{key}"}) changes_obj["new_key_2"] = MockStoreItem({"field": "new_value_for_new_key_2"}) return changes_obj + +def debug_print(*args): + """Print debug information clearly separated in the console.""" + print("\n"*2) + print("--- DEBUG ---") + for arg in args: + print("\n"*2) + print(arg) + print("\n"*2) + print("--- ----- ---") + print("\n" *2) \ No newline at end of file diff --git a/libraries/Storage/microsoft-agents-storage/tests/test_memory_storage.py b/libraries/Storage/microsoft-agents-storage/tests/test_memory_storage.py index f7b59aa1..f1d15865 100644 --- a/libraries/Storage/microsoft-agents-storage/tests/test_memory_storage.py +++ b/libraries/Storage/microsoft-agents-storage/tests/test_memory_storage.py @@ -1,22 +1,10 @@ from microsoft.agents.storage.memory_storage import MemoryStorage -from microsoft.agents.storage.storage_test_utils import CRUDStorageTests, StorageMock - - -class MemoryStorageMock(StorageMock): - - def __init__(self, initial_data: dict = None): +from microsoft.agents.storage.storage_test_utils import CRUDStorageTests +class TestMemoryStorage(CRUDStorageTests): + async def storage(self, initial_data=None): data = { key: value.store_item_to_json() for key, value in (initial_data or {}).items() } - self.storage = MemoryStorage(data) - - def get_backing_store(self): - return self.storage - - -class TestMemoryStorage(CRUDStorageTests): - - async def storage(self, initial_state=None): - return MemoryStorageMock(initial_state) + return MemoryStorage(data) \ No newline at end of file From 7638d7dd0512af6fe9b3f494d1b482fa44570a01 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Mon, 21 Jul 2025 11:38:56 -0700 Subject: [PATCH 05/13] Fixed legacy container support and added relevant tests --- .../microsoft/agents/cosmos/cosmos.py | 177 ----------- .../agents/cosmos/cosmos_db_storage.py | 131 ++++---- .../microsoft/agents/cosmos/key_ops.py | 6 + .../tests/test_cosmos_db_config.py | 241 +++++++++++++++ .../tests/test_cosmos_db_storage.py | 285 ++++++++++++++++++ .../tests/test_key_ops.py | 198 ++++++++++++ 6 files changed, 808 insertions(+), 230 deletions(-) delete mode 100644 libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos.py create mode 100644 libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py create mode 100644 libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py create mode 100644 libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos.py deleted file mode 100644 index 236be502..00000000 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos.py +++ /dev/null @@ -1,177 +0,0 @@ -# # Copyright (c) Microsoft Corporation. All rights reserved. -# # Licensed under the MIT License. - -# import json - -# from typing import TypeVar -# from threading import Lock -# from hashlib import sha256 - -# from azure.core.auth import ( -# StorageSharedKeyCredential, -# TokenCredential, -# AnonymousCredential, -# ) - -# from azure.cosmos import ( -# documents, -# http_constants, -# CosmosDict, -# ) -# from azure.cosmos.aio import ( -# ContainerProxy, -# CosmosClient, -# DatabaseProxy, -# ) -# import azure.cosmos.exceptions as cosmos_exceptions - -# from microsoft.agents.storage import Storage, StoreItem -# from microsoft.agents.storage._type_aliases import JSON -# from microsoft.agents.storage.error_handling import ignore_error, is_status_code_error - -# StoreItemT = TypeVar("StoreItemT", bound=StoreItem) - -# class CosmosDBStorage(Storage): -# """A CosmosDB based storage provider using partitioning""" - -# def __init__(self, config: CosmosDBConfig): -# """Create the storage object. - -# :param config: -# """ -# super().__init__() - -# CosmosDBStorage._validate_config(config) - -# self._config: CosmosDBConfig = config -# self._client: CosmosClient = self._create_client() -# self._database: DatabaseProxy = None -# self._container: ContainerProxy = None -# self._compatability_mode_partition_key: bool = False -# # Lock used for synchronizing container creation -# self._lock: Lock = Lock() - -# @staticmethod -# def _validate_config(config: CosmosDBConfig): -# if not config: -# raise ValueError("CosmosDBStorage: CosmosDBConfig is required.") -# if not config.cosmos_db_endpoint: -# raise ValueError("CosmosDBStorage: cosmos_db_endpoint is required.") -# if not config.auth_key: -# raise ValueError("CosmosDBStorage: auth_key is required.") -# if not config.database_id: -# raise ValueError("CosmosDBStorage: database_id is required.") -# if not config.container_id: -# raise ValueError("CosmosDBStorage: container_id is required.") - -# @staticmethod -# def _validate_suffix(config: CosmosDBConfig): - -# if config.key_suffix: -# if config.compatibility_mode: -# raise ValueError( -# "compatibilityMode cannot be true while using a keySuffix." -# ) -# suffix_escaped: str = sanitize_key(config.key_suffix) -# if suffix_escaped != config.key_suffix: -# raise ValueError( -# f"Cannot use invalid Row Key characters: {config.key_suffix} in keySuffix." -# ) - -# def _create_client(self) -> CosmosClient: - -# if self._config.url: -# if not self._config.credential: -# raise ValueError( -# "CosmosDBStorage: Credential is required when using a custom service URL." -# ) -# return CosmosClient(account_url=self._config.url, credential=self._config.credential) -# else: - -# connection_policy = self._config.cosmos_client_options.get( -# "connection_policy", documents.ConnectionPolicy() -# ) - -# # kwargs 'connection_verify' is to handle CosmosClient overwriting the -# # ConnectionPolicy.DisableSSLVerification value. -# return CosmosClient( -# self._config.cosmos_db_endpoint, -# self._config.auth_key, -# self._config.cosmos_client_options.get("consistency_level", None), -# **{ -# "connection_policy": connection_policy, -# "connection_verify": not connection_policy.DisableSSLVerification, -# }, -# ) - -# async def _read_item(key: str) -> tuple[str, StoreItemT]: -# escaped_key: str = sanitize_key( -# key, self._config.key_suffix, self._config.compatibility_mode -# ) -# read_item_response: CosmosDict = await self._container.read_item( -# escaped_key, self._get_partition_key(escaped_key) -# ) -# doc: JSON = read_item_response.get("document") -# return read_item_response["realId"], target_cls.from_json_to_store_item(doc) - -# async def _write_item(key: str, item: StoreItem) -> None: -# doc = { -# "id": sanitize_key( -# key, self._config.key_suffix, self._config.compatibility_mode -# ), -# "realId": key, -# "document": item.store_item_to_json(), -# } -# await self._container.upsert_item(body=doc) - -# async def _delete_item(self, key: str) -> None: -# escaped_key: str = sanitize_key( -# key, self._config.key_suffix, self._config.compatibility_mode -# ) -# await ignore_error(self._container.delete_item( -# escaped_key, self._get_partition_key(escaped_key) -# ), is_status_code_error(404)) - -# async def initialize(self) -> None: -# if not self._container: -# with self._lock: -# # in case another thread attempted to initialize just before acquiring the lock -# if self._container: return - -# if not self._database: -# self._database = self._client.create_database_if_not_exists( -# self._config.database_id -# ) - -# partition_key = { -# "paths": ["/id"], -# "kind": documents.PartitionKind.Hash, -# } -# try: -# self._container = self._database.create_container( -# self._config.container_id, -# partition_key, -# offer_throughput=self._config.container_throughput, -# ) -# except cosmos_exceptions.CosmosHttpResponseError as err: -# if err.status_code == http_constants.StatusCodes.CONFLICT: -# self._container = self._database.get_container_client( -# self._config.container_id -# ) -# properties = self._container.read() -# if "partitionKey" not in properties: -# self._compatability_mode_partition_key = True -# else: -# paths = properties["partitionKey"]["paths"] -# if "/partitionKey" in paths: -# self._compatability_mode_partition_key = True -# elif "/id" not in paths: -# raise Exception( -# f"Custom Partition Key Paths are not supported. {self._config.container_id} " -# "has a custom Partition Key Path of {paths[0]}." -# ) -# else: -# raise err - -# def _get_partition_key(self, key: str) -> str: -# return "" if self._compatability_mode_partition_key else key diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py index 8932e9c2..024a73be 100644 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -from typing import TypeVar +from typing import TypeVar, Union from threading import Lock from azure.cosmos import ( @@ -15,8 +15,9 @@ DatabaseProxy, ) import azure.cosmos.exceptions as cosmos_exceptions +from azure.cosmos.partition_key import NonePartitionKeyValue -from microsoft.agents.storage import Storage, StoreItem +from microsoft.agents.storage import AsyncStorageBase, StoreItem from microsoft.agents.storage._type_aliases import JSON from microsoft.agents.storage.error_handling import ignore_error, is_status_code_error @@ -25,7 +26,7 @@ StoreItemT = TypeVar("StoreItemT", bound=StoreItem) -class CosmosDBStorage(Storage): +class CosmosDBStorage(AsyncStorageBase): """A CosmosDB based storage provider using partitioning""" def __init__(self, config: CosmosDBStorageConfig): @@ -46,7 +47,6 @@ def __init__(self, config: CosmosDBStorageConfig): self._lock: Lock = Lock() def _create_client(self) -> CosmosClient: - if self._config.url: if not self._config.credential: raise ValueError( @@ -69,21 +69,44 @@ def _create_client(self) -> CosmosClient: "connection_verify": not connection_policy.DisableSSLVerification, }, ) + + def _sanitize(self, key: str) -> str: + return sanitize_key(key, self._config.key_suffix, self._config.compatibility_mode) + + async def _get_item_safe(self, raw_key: str, sanitized_key: str) -> Union[JSON, None]: + """Get item, ensuring no collision occurs. + + :param raw_key: The original key provided by the user. + :param sanitized_key: The sanitized key used for CosmosDB operations. + :return: The item if found, None otherwise. + + It's highly highly highly (highly) unlikely that a sha256 collision will ever occur. + But, just in case. + """ + + read_item_response: CosmosDict = await ignore_error(self._container.read_item( + sanitized_key, self._get_partition_key(sanitized_key) + ), lambda err: isinstance(err, cosmos_exceptions.CosmosResourceNotFoundError)) + + if (read_item_response is not None and + read_item_response.get("realId") is not None and + read_item_response.get("realId") != raw_key): + # This should never happen, but just in case... + # in fact, if this does happen, then we should probably document this + # and publish the first ever sha256 collision to the world. + raise Exception("CosmosDBStorage: Key mismatch on get_item_raw.") + + return read_item_response async def _read_item( self, key: str, *, target_cls: StoreItemT = None, **kwargs - ) -> tuple[str | None, StoreItemT | None]: + ) -> tuple[Union[str, None], Union[StoreItemT, None]]: if key == "": raise ValueError("CosmosDBStorage: Key cannot be empty.") - escaped_key: str = sanitize_key( - key, self._config.key_suffix, self._config.compatibility_mode - ) - read_item_response: CosmosDict = await ignore_error(self._container.read_item( - escaped_key, self._get_partition_key(escaped_key) - ), lambda err: isinstance(err, cosmos_exceptions.CosmosResourceNotFoundError)) - + escaped_key: str = self._sanitize(key) + read_item_response: CosmosDict = await self._get_item_safe(key, escaped_key) if read_item_response is None: return None, None @@ -91,31 +114,61 @@ async def _read_item( return read_item_response["realId"], target_cls.from_json_to_store_item(doc) async def _write_item(self, key: str, item: StoreItem) -> None: - if key == "": raise ValueError("CosmosDBStorage: Key cannot be empty.") + + escaped_key: str = self._sanitize(key) + await self._get_item_safe(key, escaped_key) # ensure no collision doc = { - "id": sanitize_key( - key, self._config.key_suffix, self._config.compatibility_mode - ), - "realId": key, + "id": escaped_key, + "realId": key, # to retrieve the raw key later "document": item.store_item_to_json(), } await self._container.upsert_item(body=doc) async def _delete_item(self, key: str) -> None: - if key == "": raise ValueError("CosmosDBStorage: Key cannot be empty.") - escaped_key: str = sanitize_key( - key, self._config.key_suffix, self._config.compatibility_mode - ) + escaped_key: str = self._sanitize(key) + await self._get_item_safe(key, escaped_key) # ensure no collision + await ignore_error(self._container.delete_item( escaped_key, self._get_partition_key(escaped_key) ), is_status_code_error(404)) + async def _create_container(self) -> None: + partition_key = { + "paths": ["/id"], + "kind": documents.PartitionKind.Hash, + } + try: + self._container = await self._database.create_container( + self._config.container_id, + partition_key, + offer_throughput=self._config.container_throughput, + ) + except cosmos_exceptions.CosmosHttpResponseError as err: + if err.status_code == http_constants.StatusCodes.CONFLICT: + self._container = self._database.get_container_client( + self._config.container_id + ) + properties = await self._container.read() + if "partitionKey" not in properties: + self._compatability_mode_partition_key = True + else: + paths = properties["partitionKey"]["paths"] + if "/_partitionKey" in paths: + self._compatability_mode_partition_key = True + elif "/id" not in paths: + raise Exception( + f"Custom Partition Key Paths are not supported. {self._config.container_id} " + "has a custom Partition Key Path of {paths[0]}." + ) + else: + raise err + async def initialize(self) -> None: if not self._container: with self._lock: @@ -127,35 +180,7 @@ async def initialize(self) -> None: self._config.database_id ) - partition_key = { - "paths": ["/id"], - "kind": documents.PartitionKind.Hash, - } - try: - self._container = await self._database.create_container( - self._config.container_id, - partition_key, - offer_throughput=self._config.container_throughput, - ) - except cosmos_exceptions.CosmosHttpResponseError as err: - if err.status_code == http_constants.StatusCodes.CONFLICT: - self._container = self._database.get_container_client( - self._config.container_id - ) - properties = await self._container.read() - if "partitionKey" not in properties: - self._compatability_mode_partition_key = True - else: - paths = properties["partitionKey"]["paths"] - if "/partitionKey" in paths: - self._compatability_mode_partition_key = True - elif "/id" not in paths: - raise Exception( - f"Custom Partition Key Paths are not supported. {self._config.container_id} " - "has a custom Partition Key Path of {paths[0]}." - ) - else: - raise err - - def _get_partition_key(self, key: str) -> str: - return "" if self._compatability_mode_partition_key else key + await self._create_container() + + def _get_partition_key(self, key: str): + return NonePartitionKeyValue if self._compatability_mode_partition_key else key diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py index e2b0a8db..4c0f3b0e 100644 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py @@ -23,12 +23,18 @@ def sanitize_key( return truncate_key(f"{key}{key_suffix}", compatibility_mode) def truncate_key(key: str, compatibility_mode: bool = True) -> str: + """ + Truncate the key to 255 characters if compatibility_mode is True. If the key is longer than 255 characters, + it will be truncated and a SHA-256 hash of the original key will be appended to minimize collisions. + """ max_key_len: int = 255 if not compatibility_mode: return key if len(key) > max_key_len: + # for now + # https://stackoverflow.com/questions/4014090/is-it-safe-to-ignore-the-possibility-of-sha-collisions-in-practice aux_hash = sha256(key.encode("utf-8")) aux_hex = aux_hash.hexdigest() diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py new file mode 100644 index 00000000..e6cdbe36 --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py @@ -0,0 +1,241 @@ +import json +import pytest + +from microsoft.agents.cosmos import CosmosDBStorageConfig + +# thank you AI, again + +@pytest.fixture() +def valid_config(): + """Fixture providing a valid CosmosDBStorageConfig for tests""" + return CosmosDBStorageConfig( + cosmos_db_endpoint="https://localhost:8081", + auth_key=( + "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGG" + "yPMbIZnqyMsEcaGQy67XIw/Jw==" + ), + database_id="test-db", + container_id="bot-storage", + ) + +@pytest.fixture() +def minimal_config(): + """Fixture providing a minimal CosmosDBStorageConfig for tests""" + return CosmosDBStorageConfig() + +@pytest.fixture() +def config_with_options(): + """Fixture providing a CosmosDBStorageConfig with all options for tests""" + return CosmosDBStorageConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db", + container_id="test_container", + cosmos_client_options={"connection_policy": "test"}, + container_throughput=800, + key_suffix="_test", + compatibility_mode=False + ) + +class TestCosmosDBStorageConfig: + + def test_constructor_with_parameters(self): + """Test creating config with direct parameters""" + config = CosmosDBStorageConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db", + container_id="test_container", + container_throughput=800, + key_suffix="_test", + compatibility_mode=False + ) + + assert config.cosmos_db_endpoint == "https://test.documents.azure.com:443/" + assert config.auth_key == "test_key" + assert config.database_id == "test_db" + assert config.container_id == "test_container" + assert config.container_throughput == 800 + assert config.key_suffix == "_test" + assert config.compatibility_mode is False + assert config.cosmos_client_options == {} + assert config.credential is None + + def test_constructor_with_defaults(self): + """Test creating config with default values""" + config = CosmosDBStorageConfig() + + assert config.cosmos_db_endpoint == "" + assert config.auth_key == "" + assert config.database_id == "" + assert config.container_id == "" + assert config.container_throughput == 400 # Default value + assert config.key_suffix == "" + assert config.compatibility_mode is False + assert config.cosmos_client_options == {} + assert config.credential is None + + def test_from_file(self, tmp_path): + """Test creating config from JSON file""" + config_file_path = tmp_path / "cosmos_config.json" + + config_data = { + "cosmos_db_endpoint": "https://localhost:8081", + "auth_key": "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==", + "database_id": "test-db", + "container_id": "bot-storage", + "container_throughput": 600, + "key_suffix": "_file", + "compatibility_mode": True, + "cosmos_client_options": {"connection_policy": "test"} + } + + with open(config_file_path, "w") as f: + json.dump(config_data, f) + + config = CosmosDBStorageConfig(filename=str(config_file_path)) + + assert config.cosmos_db_endpoint == "https://localhost:8081" + assert config.auth_key == "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==" + assert config.database_id == "test-db" + assert config.container_id == "bot-storage" + assert config.container_throughput == 600 + assert config.key_suffix == "_file" + assert config.compatibility_mode is True + assert config.cosmos_client_options == {"connection_policy": "test"} + + def test_parameter_override_file(self, tmp_path): + """Test that constructor parameters override file values""" + config_file_path = tmp_path / "cosmos_config.json" + + with open(config_file_path, "w") as f: + json.dump({ + "cosmos_db_endpoint": "https://file-endpoint.com", + "auth_key": "file_key", + "database_id": "file_db" + }, f) + + config = CosmosDBStorageConfig( + cosmos_db_endpoint="https://param-endpoint.com", + auth_key="param_key", + filename=str(config_file_path) + ) + + # Parameters should override file values + assert config.cosmos_db_endpoint == "https://param-endpoint.com" + assert config.auth_key == "param_key" + # File value should be used when parameter not provided + assert config.database_id == "file_db" + + def test_validation_success(self): + """Test successful validation with all required fields""" + config = CosmosDBStorageConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db", + container_id="test_container" + ) + + # Should not raise any exception + CosmosDBStorageConfig.validate_cosmos_db_config(config) + + def test_validation_missing_config(self): + """Test validation with None config""" + with pytest.raises(ValueError): + CosmosDBStorageConfig.validate_cosmos_db_config(None) + + def test_validation_missing_endpoint(self): + """Test validation with missing cosmos_db_endpoint""" + config = CosmosDBStorageConfig( + auth_key="test_key", + database_id="test_db", + container_id="test_container" + ) + with pytest.raises(ValueError): + CosmosDBStorageConfig.validate_cosmos_db_config(config) + + def test_validation_missing_auth_key(self): + """Test validation with missing auth_key""" + config = CosmosDBStorageConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + database_id="test_db", + container_id="test_container" + ) + with pytest.raises(ValueError): + CosmosDBStorageConfig.validate_cosmos_db_config(config) + + def test_validation_missing_database_id(self): + """Test validation with missing database_id""" + config = CosmosDBStorageConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + container_id="test_container" + ) + with pytest.raises(ValueError): + CosmosDBStorageConfig.validate_cosmos_db_config(config) + + def test_validation_missing_container_id(self): + """Test validation with missing container_id""" + config = CosmosDBStorageConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db" + ) + with pytest.raises(ValueError): + CosmosDBStorageConfig.validate_cosmos_db_config(config) + + def test_validation_suffix_with_compatibility_mode(self): + """Test validation fails when using suffix with compatibility mode""" + config = CosmosDBStorageConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db", + container_id="test_container", + key_suffix="_test", + compatibility_mode=True + ) + with pytest.raises(ValueError): + CosmosDBStorageConfig.validate_cosmos_db_config(config) + + def test_validation_invalid_suffix_characters(self): + """Test validation fails with invalid characters in suffix""" + config = CosmosDBStorageConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db", + container_id="test_container", + key_suffix="invalid/suffix\\with?bad#chars", + compatibility_mode=False + ) + with pytest.raises(ValueError, match="Cannot use invalid Row Key characters"): + CosmosDBStorageConfig.validate_cosmos_db_config(config) + + def test_validation_valid_suffix(self): + """Test validation succeeds with valid suffix""" + config = CosmosDBStorageConfig( + cosmos_db_endpoint="https://test.documents.azure.com:443/", + auth_key="test_key", + database_id="test_db", + container_id="test_container", + key_suffix="valid_suffix_123", + compatibility_mode=False + ) + # Should not raise any exception + CosmosDBStorageConfig.validate_cosmos_db_config(config) + + def test_cosmos_client_options(self): + """Test cosmos_client_options handling""" + options = {"connection_policy": "test", "consistency_level": "strong"} + config = CosmosDBStorageConfig( + cosmos_client_options=options + ) + assert config.cosmos_client_options == options + + def test_credential_parameter(self): + """Test credential parameter handling""" + # Mock credential (in real usage this would be a TokenCredential instance) + mock_credential = object() # Placeholder for actual TokenCredential + config = CosmosDBStorageConfig( + credential=mock_credential + ) + assert config.credential is mock_credential \ No newline at end of file diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py new file mode 100644 index 00000000..ab107a3d --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py @@ -0,0 +1,285 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import json, gc +from io import BytesIO + +import pytest +import pytest_asyncio + +from azure.cosmos import documents +from azure.cosmos.aio import CosmosClient +from azure.cosmos.exceptions import CosmosResourceNotFoundError + +from microsoft.agents.cosmos import CosmosDBStorage, CosmosDBStorageConfig +from microsoft.agents.cosmos.key_ops import sanitize_key + +from microsoft.agents.storage.storage_test_utils import ( + QuickCRUDStorageTests, + MockStoreItem, + MockStoreItemB, + StorageBaseline, + debug_print +) + +EMULATOR_RUNNING = True + +def create_config(compat_mode): + return CosmosDBStorageConfig( + cosmos_db_endpoint="https://localhost:8081", + auth_key=( + "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGG" + "yPMbIZnqyMsEcaGQy67XIw/Jw==" + ), + database_id="test-db", + container_id="bot-storage", + compatibility_mode=compat_mode, + container_throughput=800, + ) + +@pytest.fixture +def config(): + return create_config(compat_mode=False) + +async def create_cosmos_env(config, compat_mode=False, existing=False): + """Creates the Cosmos DB environment for testing. + + If existing is False, creates a new database and container, deleting any + existing ones with the same name. If existing is True, creates the database + and container if they do not already exist.""" + + cosmos_client = CosmosClient( + config.cosmos_db_endpoint, + config.auth_key, + ) + + if not existing: + try: + await cosmos_client.delete_database(config.database_id) + except Exception: + pass + database = await cosmos_client.create_database(id=config.database_id) + + try: + await database.delete_container(config.container_id) + except Exception: + pass + + partition_key = { + "paths": ["/_partitionKey"] if compat_mode else ["/id"], + "kind": documents.PartitionKind.Hash, + } + container_client = await database.create_container( + id=config.container_id, + partition_key=partition_key, + offer_throughput=config.container_throughput, + ) + else: + database = await cosmos_client.create_database_if_not_exists(id=config.database_id) + container_client = database.get_container_client(config.container_id) + + return container_client + +async def cosmos_db_storage_instance(compat_mode=False, existing=False): + config = create_config(compat_mode) + container_client = await create_cosmos_env(config, compat_mode=compat_mode, existing=existing) + storage = CosmosDBStorage(config) + return storage, container_client + +@pytest_asyncio.fixture() +async def cosmos_db_storage(): + storage, _ = await cosmos_db_storage_instance() + return storage + +@pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") +class dTestCosmosDBStorage(QuickCRUDStorageTests): + + def get_compat_mode(self): + return False + + async def storage(self, initial_data = None, existing=False): + storage, _ = await cosmos_db_storage_instance(compat_mode=self.get_compat_mode(), existing=existing) + if initial_data: await storage.write(initial_data) + return storage + + @pytest.mark.asyncio + async def test_initialize(self, cosmos_db_storage): + await cosmos_db_storage.initialize() + await cosmos_db_storage.initialize() + await cosmos_db_storage.write({"some_Key": MockStoreItem({"id": "123", "data": "value"})}) + await cosmos_db_storage.initialize() + assert (await cosmos_db_storage.read(["some_Key"], target_cls=MockStoreItem)) == {"some_Key": MockStoreItem({"id": "123", "data": "value"})} + + @pytest.mark.asyncio + async def test_collision_detection(self): + cosmos_storage, container_client = await cosmos_db_storage_instance() + await container_client.upsert_item({"id": "key", "realId": "fake_key", "document": {"id": "key", "value": "data"}, "partitionKey": ""}) + with pytest.raises(Exception): + await cosmos_storage.read(["key"], target_cls=MockStoreItem) + with pytest.raises(Exception): + await cosmos_storage.write({"key": MockStoreItem({"id": "item", "value": "data"})}) + with pytest.raises(Exception): + await cosmos_storage.delete(["key"]) + + @pytest.mark.asyncio + async def test_external_change_is_visible(self): + cosmos_storage, container_client = await cosmos_db_storage_instance() + assert (await cosmos_storage.read(["key"], target_cls=MockStoreItem)) == {} + assert (await cosmos_storage.read(["key2"], target_cls=MockStoreItem)) == {} + await container_client.upsert_item({"id": "key", "realId": "key", "document": {"id": "key", "value": "data"}, "partitionKey": ""}) + await container_client.upsert_item({"id": "key2", "realId": "key2", "document": {"id": "key2", "value": "new_val"}, "partitionKey": ""}) + assert (await cosmos_storage.read(["key"], target_cls=MockStoreItem))["key"] == MockStoreItem({"id": "key", "value": "data"}) + assert (await cosmos_storage.read(["key2"], target_cls=MockStoreItem))["key2"] == MockStoreItem({"id": "key2", "value": "new_val"}) + + @pytest.mark.asyncio + async def cosmos_db_storage_flow_existing_container_and_persistence(self, config, test_require_compat=False): + + config = create_config(compat_mode=test_require_compat) + container_client = await create_cosmos_env(config) + config.compatibility_mode = self.get_compat_mode() + + initial_data = { + "__some_key": MockStoreItem({"id": "item2", "value": "data2"}), + "item1"*100: MockStoreItem({"id": "item1", "value": "data1"}), + "!another_key": MockStoreItem({"id": "item3", "value": "data3"}), + "1230": MockStoreItemB({"id": "item8", "value": "data"}, False), + "key-with-dash": MockStoreItem({"id": "item4", "value": "data"}), + "key.with.dot": MockStoreItem({"id": "item5", "value": "data"}), + "key/with/slash": MockStoreItem({"id": "item6", "value": "data"}), + "another key": MockStoreItemB({"id": "item7", "value": "data"}, True), + } + + baseline_storage = StorageBaseline(initial_data) + + for key, value in initial_data.items(): + doc = { + "id": sanitize_key( + key, config.key_suffix, test_require_compat or self.get_compat_mode() + ), + "realId": key, + "document": value.store_item_to_json(), + "partitionKey": "" + } + await container_client.upsert_item(body=doc) + + storage = CosmosDBStorage(config) + storage.initialize + # debug_print(storage._compatability_mode_partition_key, self.get_compat_mode(), test_require_compat) + + assert await baseline_storage.equals(storage) + assert ( + await storage.read(["1230", "another key"], target_cls=MockStoreItemB) + ) == baseline_storage.read(["1230", "another key"]) + + changes = { + "item1"*100: MockStoreItem({"id": "item1", "value": "data1_changed"}), + "__some_key": MockStoreItem({"id": "item2", "value": "data2_changed"}), + "new_item": MockStoreItem({"id": "new_item", "value": "new_data"}), + } + + baseline_storage.write(changes) + await storage.write(changes) + + baseline_storage.delete(["!another_key", "item1"*100]) + await storage.delete(["!another_key", "item1"*100]) + assert await baseline_storage.equals(storage) + + del storage + gc.collect() + storage = CosmosDBStorage(config) + + escaped_key = storage._sanitize("item1"*100) + with pytest.raises(CosmosResourceNotFoundError): + await container_client.read_item(escaped_key, storage._get_partition_key(escaped_key)) + + escaped_key = storage._sanitize("1230") + item = (await container_client.read_item(escaped_key, storage._get_partition_key(escaped_key))).get("document") + assert ( + MockStoreItemB.from_json_to_store_item(item) + == initial_data["1230"] + ) + + @pytest.mark.asyncio + async def test_cosmos_db_storage_flow_existing_container_and_persistence(self): + await self.cosmos_db_storage_flow_existing_container_and_persistence(config) + assert True + +@pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") +class TestCosmosDBStorageWithCompat(dTestCosmosDBStorage): + def get_compat_mode(self): + return True + + @pytest.mark.asyncio + async def test_cosmos_db_storage_flow_existing_container_and_persistence(self): + await self.cosmos_db_storage_flow_existing_container_and_persistence(config) + assert True + + @pytest.mark.asyncio + async def test_cosmos_db_storage_flow_existing_container_and_persistence_with_compat(self): + await self.cosmos_db_storage_flow_existing_container_and_persistence(config, True) + assert True + +@pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") +class TestCosmosDBStorageInit: + + def test_raises_error_when_no_endpoint_provided(self, config): + config.cosmos_db_endpoint = None + with pytest.raises(ValueError): + CosmosDBStorage(config) + + def test_raises_error_when_no_auth_key_provided(self, config): + config.auth_key = None + with pytest.raises(ValueError): + CosmosDBStorage(config) + + def test_raises_error_when_suffix_provided_but_compat(self, config): + config.auth_key = None + config.compatibility_mode = True + with pytest.raises(ValueError): + CosmosDBStorage(config) + + def test_raises_error_when_no_database_id_provided(self, config): + config.database_id = None + with pytest.raises(ValueError): + CosmosDBStorage(config) + + def test_raises_error_when_no_container_id_provided(self, config): + config.container_id = None + with pytest.raises(ValueError): + CosmosDBStorage(config) + + @pytest.mark.asyncio + @pytest.mark.parametrize("compat_mode", [True, False]) + async def test_raises_error_different_partition_key(self, compat_mode): + config = create_config(compat_mode=compat_mode) + await create_cosmos_env(config, compat_mode=compat_mode) + storage = CosmosDBStorage(config) + + with pytest.raises(Exception): + + cosmos_client = CosmosClient( + config.cosmos_db_endpoint, + config.auth_key, + ) + try: + await cosmos_client.delete_database(config.database_id) + except Exception: + pass + database = await cosmos_client.create_database(id=config.database_id) + + try: + await database.delete_container(config.container_id) + except Exception: + pass + + partition_key = { + "paths": ["/fake_part_key"], + "kind": documents.PartitionKind.Hash, + } + container_client = await database.create_container( + id=config.container_id, + partition_key=partition_key, + offer_throughput=config.container_throughput, + ) + storage = CosmosDBStorage(config) + await storage.initialize() diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py b/libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py new file mode 100644 index 00000000..8e287e0c --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py @@ -0,0 +1,198 @@ +import hashlib + +from microsoft.agents.cosmos.key_ops import truncate_key, sanitize_key + +# thank you AI + +def test_sanitize_key_simple(): + """Test sanitize_key with valid keys that shouldn't change""" + assert sanitize_key("validKey123") == "validKey123" + assert sanitize_key("simple") == "simple" + assert sanitize_key("CamelCase") == "CamelCase" + assert sanitize_key("under_score") == "under_score" + assert sanitize_key("with-dash") == "with-dash" + assert sanitize_key("with.dot") == "with.dot" + +def test_sanitize_key_forbidden_chars(): + """Test sanitize_key with forbidden characters""" + # Test each forbidden character individually + assert sanitize_key("key\\value") == "key*92value" # backslash + assert sanitize_key("key?value") == "key*63value" # question mark + assert sanitize_key("key/value") == "key*47value" # forward slash + assert sanitize_key("key#value") == "key*35value" # hash + assert sanitize_key("key\tvalue") == "key*9value" # tab + assert sanitize_key("key\nvalue") == "key*10value" # newline + assert sanitize_key("key\rvalue") == "key*13value" # carriage return + assert sanitize_key("key*value") == "key*42value" # asterisk + +def test_sanitize_key_multiple_forbidden_chars(): + """Test sanitize_key with multiple forbidden characters""" + assert sanitize_key("key/with\\many?bad#chars") == "key*47with*92many*63bad*35chars" + assert sanitize_key("a\\b/c?d#e\tf\ng\rh*i") == "a*92b*47c*63d*35e*9f*10g*13h*42i" + +def test_sanitize_key_with_long_key_with_forbidden_chars(): + """Test sanitize_key with a long key that requires truncation""" + long_key = "a?2/!@\t3." * 100 # Create a long key + sanitized = sanitize_key(long_key) + assert len(sanitized) <= 255 # Should be truncated + # Ensure forbidden characters are replaced + assert "?" not in sanitized + assert "/" not in sanitized + assert "\t" not in sanitized + +def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix(): + """Test sanitize_key with a long key that requires truncation""" + long_key = "a?2/!@\t3." * 100 # Create a long key + sanitized = sanitize_key(long_key, key_suffix="_suff?#*") + assert len(sanitized) <= 255 # Should be truncated + # Ensure forbidden characters are replaced + assert "?" not in sanitized + assert "/" not in sanitized + assert "#" not in sanitized + +def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix_compat_mode(): + """Test sanitize_key with a long key that requires truncation""" + long_key = "a?2/!@\t3." * 100 # Create a long key + sanitized = sanitize_key(long_key, key_suffix="_suff?#*", compatibility_mode=True) + assert len(sanitized) <= 255 # Should be truncated + # Ensure forbidden characters are replaced + assert "?" not in sanitized + assert "/" not in sanitized + assert "#" not in sanitized + +def test_sanitize_key_empty_and_whitespace(): + """Test sanitize_key with empty and whitespace strings""" + assert sanitize_key("") == "" + assert sanitize_key(" ") == " " # spaces are allowed + assert sanitize_key(" key ") == " key " # leading/trailing spaces preserved + +def test_sanitize_key_with_suffix(): + """Test sanitize_key with key_suffix parameter""" + assert sanitize_key("key", key_suffix="_suffix") == "key_suffix" + assert sanitize_key("test", key_suffix="123") == "test123" + assert sanitize_key("key/value", key_suffix="_clean") == "key*47value_clean" + assert sanitize_key("", key_suffix="_suffix") == "_suffix" + +def test_sanitize_key_suffix_with_truncation(): + """Test sanitize_key with suffix that causes truncation""" + long_key = "a" * 250 + suffix = "_suffix" + result = sanitize_key(long_key, key_suffix=suffix, compatibility_mode=True) + assert len(result) <= 255 + assert result.endswith(suffix) or len(result) == 255 # Either has suffix or was truncated + +def test_sanitize_key_truncation_compatibility_mode(): + """Test sanitize_key with truncation in compatibility mode""" + long_key = "a" * 300 + result = sanitize_key(long_key, compatibility_mode=True) + assert len(result) <= 255 + + # Should contain hash when truncated + very_long_key = "b" * 500 + result2 = sanitize_key(very_long_key, compatibility_mode=True) + assert len(result2) == 255 + +def test_sanitize_key_no_truncation(): + """Test sanitize_key without compatibility mode (no truncation)""" + long_key = "a" * 300 + result = sanitize_key(long_key, compatibility_mode=False) + assert result == long_key # Should be unchanged + assert len(result) == 300 + +def test_truncate_key_short_strings(): + """Test truncate_key with strings shorter than 255 characters""" + assert truncate_key("short") == "short" + assert truncate_key("a" * 254) == "a" * 254 + assert truncate_key("a" * 255) == "a" * 255 + +def test_truncate_key_long_strings(): + """Test truncate_key with strings longer than 255 characters""" + long_key = "a" * 300 + result = truncate_key(long_key) + assert len(result) == 255 + + # Result should end with SHA256 hash + expected_hash = hashlib.sha256(long_key.encode("utf-8")).hexdigest() + assert result.endswith(expected_hash) + + # First part should be original key truncated + expected_prefix_len = 255 - len(expected_hash) + assert result.startswith("a" * expected_prefix_len) + +def test_truncate_key_compatibility_mode_disabled(): + """Test truncate_key with compatibility_mode=False""" + long_key = "a" * 300 + assert truncate_key(long_key, compatibility_mode=False) == long_key + + very_long_key = "x" * 1000 + assert truncate_key(very_long_key, compatibility_mode=False) == very_long_key + + # Should work with any characters when compatibility mode is off + special_chars = "key/with\\special?chars#and\ttabs\nand\rmore*" + assert truncate_key(special_chars, compatibility_mode=False) == special_chars + +def test_truncate_key_exactly_255_chars(): + """Test truncate_key with exactly 255 characters""" + key_255 = "a" * 255 + assert truncate_key(key_255) == key_255 + assert len(truncate_key(key_255)) == 255 + +def test_truncate_key_256_chars(): + """Test truncate_key with 256 characters (one over limit)""" + key_256 = "a" * 256 + result = truncate_key(key_256) + assert len(result) == 255 + assert result != key_256 + +def test_truncate_key_hash_consistency(): + """Test that truncate_key produces consistent hashes for same input""" + long_key = "consistent_test_key_" * 20 # > 255 chars + result1 = truncate_key(long_key) + result2 = truncate_key(long_key) + assert result1 == result2 + assert len(result1) == 255 + +def test_truncate_key_different_inputs_different_outputs(): + """Test that different long keys produce different truncated results""" + key1 = "a" * 300 + key2 = "b" * 300 + result1 = truncate_key(key1) + result2 = truncate_key(key2) + assert result1 != result2 + assert len(result1) == len(result2) == 255 + +def test_sanitize_key_integration(): + """Integration test combining forbidden chars, suffix, and truncation""" + # Key with forbidden chars that will be long after sanitization + suffix + base_key = "test/key\\with?many#forbidden\tchars\nand\rmore*" * 10 + suffix = "_integration_test" + + result = sanitize_key(base_key, key_suffix=suffix, compatibility_mode=True) + + # Should be sanitized and truncated + assert len(result) <= 255 + assert "*47" in result or "*92" in result # Contains sanitized chars + + # Test without truncation + result_no_trunc = sanitize_key(base_key, key_suffix=suffix, compatibility_mode=False) + assert "*47" in result_no_trunc or "*92" in result_no_trunc # Contains sanitized chars + assert result_no_trunc.endswith(suffix) + +def test_edge_cases(): + """Test various edge cases""" + # Unicode characters (should be preserved) + assert sanitize_key("key_ñ_测试") == "key_ñ_测试" + + # Numbers only + assert sanitize_key("123456789") == "123456789" + + # Mixed case with special chars + result = sanitize_key("MyKey/WithSlash") + assert result == "MyKey*47WithSlash" + + # Key that becomes exactly 255 chars after sanitization + # Create a key that will be exactly 255 after replacing one char + base = "a" * 252 + "/" # 253 chars, "/" becomes "*47" (3 chars) = 255 total + result = sanitize_key(base) + assert len(result) == 255 + assert result.endswith("*47") \ No newline at end of file From 50dcff9023b7f8b74c00b771f05281f15330182c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Mon, 21 Jul 2025 11:40:23 -0700 Subject: [PATCH 06/13] Revised AsyncStorageBase export/imports --- .../microsoft/agents/blob/blob_storage.py | 3 +-- .../Storage/microsoft-agents-blob/tests/test_blob_storage.py | 1 - .../microsoft/agents/storage/__init__.py | 4 ++-- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/libraries/Storage/microsoft-agents-blob/microsoft/agents/blob/blob_storage.py b/libraries/Storage/microsoft-agents-blob/microsoft/agents/blob/blob_storage.py index 39c54a52..2f8e0178 100644 --- a/libraries/Storage/microsoft-agents-blob/microsoft/agents/blob/blob_storage.py +++ b/libraries/Storage/microsoft-agents-blob/microsoft/agents/blob/blob_storage.py @@ -7,8 +7,7 @@ BlobServiceClient, ) -from microsoft.agents.storage import StoreItem -from microsoft.agents.storage.storage import AsyncStorageBase +from microsoft.agents.storage import AsyncStorageBase, StoreItem from microsoft.agents.storage._type_aliases import JSON from microsoft.agents.storage.error_handling import ignore_error, is_status_code_error diff --git a/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py b/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py index cba9f88d..616148c6 100644 --- a/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py +++ b/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py @@ -180,5 +180,4 @@ async def test_blob_storage_flow_existing_container_and_persistence(self): == initial_data["1230"] ) - # teardown await container_client.delete_container() diff --git a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/__init__.py b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/__init__.py index be2e6079..1d54743e 100644 --- a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/__init__.py +++ b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/__init__.py @@ -1,5 +1,5 @@ from .store_item import StoreItem -from .storage import Storage +from .storage import Storage, AsyncStorageBase from .memory_storage import MemoryStorage -__all__ = ["StoreItem", "Storage", "MemoryStorage"] +__all__ = ["StoreItem", "Storage", "AsyncStorageBase", "MemoryStorage"] From 3e90218106438c9d21df03a0068d81c4f6b96b48 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Mon, 21 Jul 2025 12:44:45 -0700 Subject: [PATCH 07/13] Fixed final unit tests and refined documentation --- .../tests/test_blob_storage.py | 29 ++- .../agents/cosmos/cosmos_db_storage.py | 91 ++++--- .../agents/cosmos/cosmos_db_storage_config.py | 10 +- .../microsoft/agents/cosmos/key_ops.py | 4 +- .../microsoft-agents-cosmos/pyproject.toml | 23 ++ .../tests/test_cosmos_db_config.py | 64 ++--- .../tests/test_cosmos_db_storage.py | 230 ++++++++++-------- .../tests/test_key_ops.py | 56 +++-- .../agents/storage/error_handling.py | 2 +- .../agents/storage/storage_test_utils.py | 33 ++- .../tests/test_memory_storage.py | 3 +- 11 files changed, 344 insertions(+), 201 deletions(-) create mode 100644 libraries/Storage/microsoft-agents-cosmos/pyproject.toml diff --git a/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py b/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py index 616148c6..468c4274 100644 --- a/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py +++ b/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py @@ -19,6 +19,7 @@ EMULATOR_RUNNING = False + async def blob_storage_instance(existing=False): # Default Azure Storage Emulator connection string @@ -54,6 +55,7 @@ async def blob_storage_instance(existing=False): storage = BlobStorage(blob_storage_config) return storage, container_client + @pytest_asyncio.fixture async def blob_storage(): @@ -65,6 +67,7 @@ async def blob_storage(): # teardown await container_client.delete_container() + @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") class TestBlobStorage(CRUDStorageTests): @@ -94,12 +97,26 @@ async def test_initialize(self, blob_storage): @pytest.mark.asyncio async def test_external_change_is_visible(self): blob_storage, container_client = await blob_storage_instance() - assert (await blob_storage.read(["key"], target_cls=MockStoreItem)) == MockStoreItem({"id": "item", "value": "data"}) - assert (await blob_storage.read(["key2"], target_cls=MockStoreItem)) == MockStoreItem({"id": "another_item", "value": "another_value"}) - await container_client.upload_blob(name="key", data=json.dumps({"id": "item", "value": "data"}), overwrite=True) - await container_client.upload_blob(name="key2", data=json.dumps({"id": "another_item", "value": "new_val"}), overwrite=True) - assert (await blob_storage.read(["key"], target_cls=MockStoreItem)) == MockStoreItem({"id": "item", "value": "data"}) - assert (await blob_storage.read(["key2"], target_cls=MockStoreItem)) == MockStoreItem({"id": "another_item", "value": "new_val"}) + assert ( + await blob_storage.read(["key"], target_cls=MockStoreItem) + ) == MockStoreItem({"id": "item", "value": "data"}) + assert ( + await blob_storage.read(["key2"], target_cls=MockStoreItem) + ) == MockStoreItem({"id": "another_item", "value": "another_value"}) + await container_client.upload_blob( + name="key", data=json.dumps({"id": "item", "value": "data"}), overwrite=True + ) + await container_client.upload_blob( + name="key2", + data=json.dumps({"id": "another_item", "value": "new_val"}), + overwrite=True, + ) + assert ( + await blob_storage.read(["key"], target_cls=MockStoreItem) + ) == MockStoreItem({"id": "item", "value": "data"}) + assert ( + await blob_storage.read(["key2"], target_cls=MockStoreItem) + ) == MockStoreItem({"id": "another_item", "value": "new_val"}) @pytest.mark.asyncio async def test_blob_storage_flow_existing_container_and_persistence(self): diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py index 024a73be..aa6a2bd2 100644 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py @@ -26,6 +26,11 @@ StoreItemT = TypeVar("StoreItemT", bound=StoreItem) +cosmos_resource_not_found = lambda err: isinstance( + err, cosmos_exceptions.CosmosResourceNotFoundError +) + + class CosmosDBStorage(AsyncStorageBase): """A CosmosDB based storage provider using partitioning""" @@ -45,14 +50,16 @@ def __init__(self, config: CosmosDBStorageConfig): self._compatability_mode_partition_key: bool = False # Lock used for synchronizing container creation self._lock: Lock = Lock() - + def _create_client(self) -> CosmosClient: if self._config.url: if not self._config.credential: raise ValueError( "CosmosDBStorage: Credential is required when using a custom service URL." ) - return CosmosClient(account_url=self._config.url, credential=self._config.credential) + return CosmosClient( + account_url=self._config.url, credential=self._config.credential + ) connection_policy = self._config.cosmos_client_options.get( "connection_policy", documents.ConnectionPolicy() @@ -63,19 +70,25 @@ def _create_client(self) -> CosmosClient: return CosmosClient( self._config.cosmos_db_endpoint, self._config.auth_key, - consistency_level=self._config.cosmos_client_options.get("consistency_level", None), + consistency_level=self._config.cosmos_client_options.get( + "consistency_level", None + ), **{ "connection_policy": connection_policy, "connection_verify": not connection_policy.DisableSSLVerification, }, ) - + def _sanitize(self, key: str) -> str: - return sanitize_key(key, self._config.key_suffix, self._config.compatibility_mode) - - async def _get_item_safe(self, raw_key: str, sanitized_key: str) -> Union[JSON, None]: + return sanitize_key( + key, self._config.key_suffix, self._config.compatibility_mode + ) + + async def _get_item_safe( + self, raw_key: str, sanitized_key: str + ) -> Union[JSON, None]: """Get item, ensuring no collision occurs. - + :param raw_key: The original key provided by the user. :param sanitized_key: The sanitized key used for CosmosDB operations. :return: The item if found, None otherwise. @@ -84,24 +97,29 @@ async def _get_item_safe(self, raw_key: str, sanitized_key: str) -> Union[JSON, But, just in case. """ - read_item_response: CosmosDict = await ignore_error(self._container.read_item( - sanitized_key, self._get_partition_key(sanitized_key) - ), lambda err: isinstance(err, cosmos_exceptions.CosmosResourceNotFoundError)) + read_item_response: CosmosDict = await ignore_error( + self._container.read_item( + sanitized_key, self._get_partition_key(sanitized_key) + ), + cosmos_resource_not_found, + ) - if (read_item_response is not None and - read_item_response.get("realId") is not None and - read_item_response.get("realId") != raw_key): + if ( + read_item_response is not None + and read_item_response.get("realId") is not None + and read_item_response.get("realId") != raw_key + ): # This should never happen, but just in case... # in fact, if this does happen, then we should probably document this # and publish the first ever sha256 collision to the world. raise Exception("CosmosDBStorage: Key mismatch on get_item_raw.") - + return read_item_response async def _read_item( self, key: str, *, target_cls: StoreItemT = None, **kwargs ) -> tuple[Union[str, None], Union[StoreItemT, None]]: - + if key == "": raise ValueError("CosmosDBStorage: Key cannot be empty.") @@ -109,20 +127,20 @@ async def _read_item( read_item_response: CosmosDict = await self._get_item_safe(key, escaped_key) if read_item_response is None: return None, None - + doc: JSON = read_item_response.get("document") return read_item_response["realId"], target_cls.from_json_to_store_item(doc) async def _write_item(self, key: str, item: StoreItem) -> None: if key == "": raise ValueError("CosmosDBStorage: Key cannot be empty.") - + escaped_key: str = self._sanitize(key) await self._get_item_safe(key, escaped_key) # ensure no collision doc = { "id": escaped_key, - "realId": key, # to retrieve the raw key later + "realId": key, # to retrieve the raw key later "document": item.store_item_to_json(), } await self._container.upsert_item(body=doc) @@ -132,11 +150,14 @@ async def _delete_item(self, key: str) -> None: raise ValueError("CosmosDBStorage: Key cannot be empty.") escaped_key: str = self._sanitize(key) - await self._get_item_safe(key, escaped_key) # ensure no collision + await self._get_item_safe(key, escaped_key) # ensure no collision - await ignore_error(self._container.delete_item( - escaped_key, self._get_partition_key(escaped_key) - ), is_status_code_error(404)) + await ignore_error( + self._container.delete_item( + escaped_key, self._get_partition_key(escaped_key) + ), + cosmos_resource_not_found, + ) async def _create_container(self) -> None: partition_key = { @@ -155,17 +176,18 @@ async def _create_container(self) -> None: self._config.container_id ) properties = await self._container.read() - if "partitionKey" not in properties: + # if "partitionKey" not in properties: + # self._compatability_mode_partition_key = True + # else: + # containers created had no partition key, so the default was "/_partitionKey" + paths = properties["partitionKey"]["paths"] + if "/_partitionKey" in paths: self._compatability_mode_partition_key = True - else: - paths = properties["partitionKey"]["paths"] - if "/_partitionKey" in paths: - self._compatability_mode_partition_key = True - elif "/id" not in paths: - raise Exception( - f"Custom Partition Key Paths are not supported. {self._config.container_id} " - "has a custom Partition Key Path of {paths[0]}." - ) + elif "/id" not in paths: + raise Exception( + f"Custom Partition Key Paths are not supported. {self._config.container_id} " + "has a custom Partition Key Path of {paths[0]}." + ) else: raise err @@ -173,7 +195,8 @@ async def initialize(self) -> None: if not self._container: with self._lock: # in case another thread attempted to initialize just before acquiring the lock - if self._container: return + if self._container: + return if not self._database: self._database = await self._client.create_database_if_not_exists( diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py index 2042887b..1ebddb5f 100644 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py @@ -5,6 +5,7 @@ from .key_ops import sanitize_key + class CosmosDBStorageConfig: """The class for partitioned CosmosDB configuration for the Azure Bot Framework.""" @@ -35,6 +36,8 @@ def __init__( key characters. (e.g. not: '\\', '?', '/', '#', '*') :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb max key length of 255. + :param url: The URL to the CosmosDB resource. + :param credential: The TokenCredential to use for authentication. :return CosmosDBConfig: """ config_file: str = kwargs.get("filename", "") @@ -61,6 +64,9 @@ def __init__( @staticmethod def validate_cosmos_db_config(config: "CosmosDBStorageConfig") -> None: + """Validate the CosmosDBConfig object. + + This is used prior to the creation of the CosmosDBStorage object.""" if not config: raise ValueError("CosmosDBStorage: CosmosDBConfig is required.") if not config.cosmos_db_endpoint: @@ -71,7 +77,7 @@ def validate_cosmos_db_config(config: "CosmosDBStorageConfig") -> None: raise ValueError("CosmosDBStorage: database_id is required.") if not config.container_id: raise ValueError("CosmosDBStorage: container_id is required.") - + CosmosDBStorageConfig._validate_suffix(config) @staticmethod @@ -85,4 +91,4 @@ def _validate_suffix(config: "CosmosDBStorageConfig") -> None: if suffix_escaped != config.key_suffix: raise ValueError( f"Cannot use invalid Row Key characters: {config.key_suffix} in keySuffix." - ) \ No newline at end of file + ) diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py index 4c0f3b0e..1d6941f9 100644 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py +++ b/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py @@ -1,5 +1,6 @@ from hashlib import sha256 + def sanitize_key( key: str, key_suffix: str = "", compatibility_mode: bool = True ) -> str: @@ -22,6 +23,7 @@ def sanitize_key( key = "".join(map(lambda x: "*" + str(ord(x)) if x in bad_chars else x, key)) return truncate_key(f"{key}{key_suffix}", compatibility_mode) + def truncate_key(key: str, compatibility_mode: bool = True) -> str: """ Truncate the key to 255 characters if compatibility_mode is True. If the key is longer than 255 characters, @@ -40,4 +42,4 @@ def truncate_key(key: str, compatibility_mode: bool = True) -> str: key = key[0 : max_key_len - len(aux_hex)] + aux_hex - return key \ No newline at end of file + return key diff --git a/libraries/Storage/microsoft-agents-cosmos/pyproject.toml b/libraries/Storage/microsoft-agents-cosmos/pyproject.toml new file mode 100644 index 00000000..152df9f4 --- /dev/null +++ b/libraries/Storage/microsoft-agents-cosmos/pyproject.toml @@ -0,0 +1,23 @@ +[build-system] +requires = ["setuptools"] +build-backend = "setuptools.build_meta" + +[project] +name = "microsoft-agents-cosmos" +version = "0.0.0a1" +description = "A Cosmos DB storage library for Microsoft Agents" +authors = [{name = "Microsoft Corporation"}] +requires-python = ">=3.9" +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", +] +dependencies = [ + "microsoft.agents.storage", + "azure-core", + "azure-cosmos", +] + +[project.urls] +"Homepage" = "https://github.com/microsoft/microsoft-agents-protocol" diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py index e6cdbe36..81f79202 100644 --- a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py +++ b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py @@ -5,6 +5,7 @@ # thank you AI, again + @pytest.fixture() def valid_config(): """Fixture providing a valid CosmosDBStorageConfig for tests""" @@ -18,11 +19,13 @@ def valid_config(): container_id="bot-storage", ) + @pytest.fixture() def minimal_config(): """Fixture providing a minimal CosmosDBStorageConfig for tests""" return CosmosDBStorageConfig() + @pytest.fixture() def config_with_options(): """Fixture providing a CosmosDBStorageConfig with all options for tests""" @@ -34,9 +37,10 @@ def config_with_options(): cosmos_client_options={"connection_policy": "test"}, container_throughput=800, key_suffix="_test", - compatibility_mode=False + compatibility_mode=False, ) + class TestCosmosDBStorageConfig: def test_constructor_with_parameters(self): @@ -48,9 +52,9 @@ def test_constructor_with_parameters(self): container_id="test_container", container_throughput=800, key_suffix="_test", - compatibility_mode=False + compatibility_mode=False, ) - + assert config.cosmos_db_endpoint == "https://test.documents.azure.com:443/" assert config.auth_key == "test_key" assert config.database_id == "test_db" @@ -64,7 +68,7 @@ def test_constructor_with_parameters(self): def test_constructor_with_defaults(self): """Test creating config with default values""" config = CosmosDBStorageConfig() - + assert config.cosmos_db_endpoint == "" assert config.auth_key == "" assert config.database_id == "" @@ -87,7 +91,7 @@ def test_from_file(self, tmp_path): "container_throughput": 600, "key_suffix": "_file", "compatibility_mode": True, - "cosmos_client_options": {"connection_policy": "test"} + "cosmos_client_options": {"connection_policy": "test"}, } with open(config_file_path, "w") as f: @@ -96,7 +100,10 @@ def test_from_file(self, tmp_path): config = CosmosDBStorageConfig(filename=str(config_file_path)) assert config.cosmos_db_endpoint == "https://localhost:8081" - assert config.auth_key == "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==" + assert ( + config.auth_key + == "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==" + ) assert config.database_id == "test-db" assert config.container_id == "bot-storage" assert config.container_throughput == 600 @@ -109,16 +116,19 @@ def test_parameter_override_file(self, tmp_path): config_file_path = tmp_path / "cosmos_config.json" with open(config_file_path, "w") as f: - json.dump({ - "cosmos_db_endpoint": "https://file-endpoint.com", - "auth_key": "file_key", - "database_id": "file_db" - }, f) + json.dump( + { + "cosmos_db_endpoint": "https://file-endpoint.com", + "auth_key": "file_key", + "database_id": "file_db", + }, + f, + ) config = CosmosDBStorageConfig( cosmos_db_endpoint="https://param-endpoint.com", auth_key="param_key", - filename=str(config_file_path) + filename=str(config_file_path), ) # Parameters should override file values @@ -133,9 +143,9 @@ def test_validation_success(self): cosmos_db_endpoint="https://test.documents.azure.com:443/", auth_key="test_key", database_id="test_db", - container_id="test_container" + container_id="test_container", ) - + # Should not raise any exception CosmosDBStorageConfig.validate_cosmos_db_config(config) @@ -147,9 +157,7 @@ def test_validation_missing_config(self): def test_validation_missing_endpoint(self): """Test validation with missing cosmos_db_endpoint""" config = CosmosDBStorageConfig( - auth_key="test_key", - database_id="test_db", - container_id="test_container" + auth_key="test_key", database_id="test_db", container_id="test_container" ) with pytest.raises(ValueError): CosmosDBStorageConfig.validate_cosmos_db_config(config) @@ -159,7 +167,7 @@ def test_validation_missing_auth_key(self): config = CosmosDBStorageConfig( cosmos_db_endpoint="https://test.documents.azure.com:443/", database_id="test_db", - container_id="test_container" + container_id="test_container", ) with pytest.raises(ValueError): CosmosDBStorageConfig.validate_cosmos_db_config(config) @@ -169,7 +177,7 @@ def test_validation_missing_database_id(self): config = CosmosDBStorageConfig( cosmos_db_endpoint="https://test.documents.azure.com:443/", auth_key="test_key", - container_id="test_container" + container_id="test_container", ) with pytest.raises(ValueError): CosmosDBStorageConfig.validate_cosmos_db_config(config) @@ -179,7 +187,7 @@ def test_validation_missing_container_id(self): config = CosmosDBStorageConfig( cosmos_db_endpoint="https://test.documents.azure.com:443/", auth_key="test_key", - database_id="test_db" + database_id="test_db", ) with pytest.raises(ValueError): CosmosDBStorageConfig.validate_cosmos_db_config(config) @@ -192,7 +200,7 @@ def test_validation_suffix_with_compatibility_mode(self): database_id="test_db", container_id="test_container", key_suffix="_test", - compatibility_mode=True + compatibility_mode=True, ) with pytest.raises(ValueError): CosmosDBStorageConfig.validate_cosmos_db_config(config) @@ -205,7 +213,7 @@ def test_validation_invalid_suffix_characters(self): database_id="test_db", container_id="test_container", key_suffix="invalid/suffix\\with?bad#chars", - compatibility_mode=False + compatibility_mode=False, ) with pytest.raises(ValueError, match="Cannot use invalid Row Key characters"): CosmosDBStorageConfig.validate_cosmos_db_config(config) @@ -218,7 +226,7 @@ def test_validation_valid_suffix(self): database_id="test_db", container_id="test_container", key_suffix="valid_suffix_123", - compatibility_mode=False + compatibility_mode=False, ) # Should not raise any exception CosmosDBStorageConfig.validate_cosmos_db_config(config) @@ -226,16 +234,12 @@ def test_validation_valid_suffix(self): def test_cosmos_client_options(self): """Test cosmos_client_options handling""" options = {"connection_policy": "test", "consistency_level": "strong"} - config = CosmosDBStorageConfig( - cosmos_client_options=options - ) + config = CosmosDBStorageConfig(cosmos_client_options=options) assert config.cosmos_client_options == options def test_credential_parameter(self): """Test credential parameter handling""" # Mock credential (in real usage this would be a TokenCredential instance) mock_credential = object() # Placeholder for actual TokenCredential - config = CosmosDBStorageConfig( - credential=mock_credential - ) - assert config.credential is mock_credential \ No newline at end of file + config = CosmosDBStorageConfig(credential=mock_credential) + assert config.credential is mock_credential diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py index ab107a3d..b004a69f 100644 --- a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py +++ b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py @@ -19,10 +19,10 @@ MockStoreItem, MockStoreItemB, StorageBaseline, - debug_print ) -EMULATOR_RUNNING = True +EMULATOR_RUNNING = False + def create_config(compat_mode): return CosmosDBStorageConfig( @@ -37,13 +37,15 @@ def create_config(compat_mode): container_throughput=800, ) + @pytest.fixture def config(): return create_config(compat_mode=False) + async def create_cosmos_env(config, compat_mode=False, existing=False): """Creates the Cosmos DB environment for testing. - + If existing is False, creates a new database and container, deleting any existing ones with the same name. If existing is True, creates the database and container if they do not already exist.""" @@ -75,49 +77,144 @@ async def create_cosmos_env(config, compat_mode=False, existing=False): offer_throughput=config.container_throughput, ) else: - database = await cosmos_client.create_database_if_not_exists(id=config.database_id) + database = await cosmos_client.create_database_if_not_exists( + id=config.database_id + ) container_client = database.get_container_client(config.container_id) return container_client + async def cosmos_db_storage_instance(compat_mode=False, existing=False): config = create_config(compat_mode) - container_client = await create_cosmos_env(config, compat_mode=compat_mode, existing=existing) + container_client = await create_cosmos_env( + config, compat_mode=compat_mode, existing=existing + ) storage = CosmosDBStorage(config) return storage, container_client + @pytest_asyncio.fixture() async def cosmos_db_storage(): storage, _ = await cosmos_db_storage_instance() return storage + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_require_compat", [True, False]) +async def test_cosmos_db_storage_flow_existing_container_and_persistence( + test_require_compat, +): + + config = create_config(compat_mode=test_require_compat) + container_client = await create_cosmos_env(config) + + initial_data = { + "__some_key": MockStoreItem({"id": "item2", "value": "data2"}), + "?test": MockStoreItem({"id": "?test", "value": "data1"}), + "!another_key": MockStoreItem({"id": "item3", "value": "data3"}), + "1230": MockStoreItemB({"id": "item8", "value": "data"}, False), + "key-with-dash": MockStoreItem({"id": "item4", "value": "data"}), + "key.with.dot": MockStoreItem({"id": "item5", "value": "data"}), + "key/with/slash": MockStoreItem({"id": "item6", "value": "data"}), + "another key": MockStoreItemB({"id": "item7", "value": "data"}, True), + } + + baseline_storage = StorageBaseline(initial_data) + + for key, value in initial_data.items(): + doc = { + "id": sanitize_key( + key, + config.key_suffix, + test_require_compat, + ), + "realId": key, + "document": value.store_item_to_json(), + } + await container_client.upsert_item(body=doc) + + storage = CosmosDBStorage(config) + assert await baseline_storage.equals(storage) + assert ( + await storage.read(["1230", "another key"], target_cls=MockStoreItemB) + ) == baseline_storage.read(["1230", "another key"]) + + changes = { + "?test": MockStoreItem({"id": "?test", "value": "data1_changed"}), + "__some_key": MockStoreItem({"id": "item2", "value": "data2_changed"}), + "new_item": MockStoreItem({"id": "new_item", "value": "new_data"}), + } + + baseline_storage.write(changes) + await storage.write(changes) + + baseline_storage.delete(["!another_key", "?test"]) + await storage.delete(["!another_key", "?test"]) + assert await baseline_storage.equals(storage) + + del storage + gc.collect() + storage = CosmosDBStorage(config) + + escaped_key = storage._sanitize("?test") + with pytest.raises(CosmosResourceNotFoundError): + await container_client.read_item( + escaped_key, storage._get_partition_key(escaped_key) + ) + + escaped_key = storage._sanitize("1230") + item = ( + await container_client.read_item( + escaped_key, storage._get_partition_key(escaped_key) + ) + ).get("document") + assert MockStoreItemB.from_json_to_store_item(item) == initial_data["1230"] + + @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") -class dTestCosmosDBStorage(QuickCRUDStorageTests): +class TestCosmosDBStorage(QuickCRUDStorageTests): def get_compat_mode(self): return False - - async def storage(self, initial_data = None, existing=False): - storage, _ = await cosmos_db_storage_instance(compat_mode=self.get_compat_mode(), existing=existing) - if initial_data: await storage.write(initial_data) + + async def storage(self, initial_data=None, existing=False): + storage, _ = await cosmos_db_storage_instance( + compat_mode=self.get_compat_mode(), existing=existing + ) + if initial_data: + await storage.write(initial_data) return storage @pytest.mark.asyncio async def test_initialize(self, cosmos_db_storage): await cosmos_db_storage.initialize() await cosmos_db_storage.initialize() - await cosmos_db_storage.write({"some_Key": MockStoreItem({"id": "123", "data": "value"})}) + await cosmos_db_storage.write( + {"some_Key": MockStoreItem({"id": "123", "data": "value"})} + ) await cosmos_db_storage.initialize() - assert (await cosmos_db_storage.read(["some_Key"], target_cls=MockStoreItem)) == {"some_Key": MockStoreItem({"id": "123", "data": "value"})} + assert ( + await cosmos_db_storage.read(["some_Key"], target_cls=MockStoreItem) + ) == {"some_Key": MockStoreItem({"id": "123", "data": "value"})} @pytest.mark.asyncio async def test_collision_detection(self): cosmos_storage, container_client = await cosmos_db_storage_instance() - await container_client.upsert_item({"id": "key", "realId": "fake_key", "document": {"id": "key", "value": "data"}, "partitionKey": ""}) + await container_client.upsert_item( + { + "id": "key", + "realId": "fake_key", + "document": {"id": "key", "value": "data"}, + "partitionKey": "", + } + ) with pytest.raises(Exception): await cosmos_storage.read(["key"], target_cls=MockStoreItem) with pytest.raises(Exception): - await cosmos_storage.write({"key": MockStoreItem({"id": "item", "value": "data"})}) + await cosmos_storage.write( + {"key": MockStoreItem({"id": "item", "value": "data"})} + ) with pytest.raises(Exception): await cosmos_storage.delete(["key"]) @@ -126,98 +223,35 @@ async def test_external_change_is_visible(self): cosmos_storage, container_client = await cosmos_db_storage_instance() assert (await cosmos_storage.read(["key"], target_cls=MockStoreItem)) == {} assert (await cosmos_storage.read(["key2"], target_cls=MockStoreItem)) == {} - await container_client.upsert_item({"id": "key", "realId": "key", "document": {"id": "key", "value": "data"}, "partitionKey": ""}) - await container_client.upsert_item({"id": "key2", "realId": "key2", "document": {"id": "key2", "value": "new_val"}, "partitionKey": ""}) - assert (await cosmos_storage.read(["key"], target_cls=MockStoreItem))["key"] == MockStoreItem({"id": "key", "value": "data"}) - assert (await cosmos_storage.read(["key2"], target_cls=MockStoreItem))["key2"] == MockStoreItem({"id": "key2", "value": "new_val"}) - - @pytest.mark.asyncio - async def cosmos_db_storage_flow_existing_container_and_persistence(self, config, test_require_compat=False): - - config = create_config(compat_mode=test_require_compat) - container_client = await create_cosmos_env(config) - config.compatibility_mode = self.get_compat_mode() - - initial_data = { - "__some_key": MockStoreItem({"id": "item2", "value": "data2"}), - "item1"*100: MockStoreItem({"id": "item1", "value": "data1"}), - "!another_key": MockStoreItem({"id": "item3", "value": "data3"}), - "1230": MockStoreItemB({"id": "item8", "value": "data"}, False), - "key-with-dash": MockStoreItem({"id": "item4", "value": "data"}), - "key.with.dot": MockStoreItem({"id": "item5", "value": "data"}), - "key/with/slash": MockStoreItem({"id": "item6", "value": "data"}), - "another key": MockStoreItemB({"id": "item7", "value": "data"}, True), - } - - baseline_storage = StorageBaseline(initial_data) - - for key, value in initial_data.items(): - doc = { - "id": sanitize_key( - key, config.key_suffix, test_require_compat or self.get_compat_mode() - ), - "realId": key, - "document": value.store_item_to_json(), - "partitionKey": "" + await container_client.upsert_item( + { + "id": "key", + "realId": "key", + "document": {"id": "key", "value": "data"}, + "partitionKey": "", + } + ) + await container_client.upsert_item( + { + "id": "key2", + "realId": "key2", + "document": {"id": "key2", "value": "new_val"}, + "partitionKey": "", } - await container_client.upsert_item(body=doc) - - storage = CosmosDBStorage(config) - storage.initialize - # debug_print(storage._compatability_mode_partition_key, self.get_compat_mode(), test_require_compat) - - assert await baseline_storage.equals(storage) - assert ( - await storage.read(["1230", "another key"], target_cls=MockStoreItemB) - ) == baseline_storage.read(["1230", "another key"]) - - changes = { - "item1"*100: MockStoreItem({"id": "item1", "value": "data1_changed"}), - "__some_key": MockStoreItem({"id": "item2", "value": "data2_changed"}), - "new_item": MockStoreItem({"id": "new_item", "value": "new_data"}), - } - - baseline_storage.write(changes) - await storage.write(changes) - - baseline_storage.delete(["!another_key", "item1"*100]) - await storage.delete(["!another_key", "item1"*100]) - assert await baseline_storage.equals(storage) - - del storage - gc.collect() - storage = CosmosDBStorage(config) - - escaped_key = storage._sanitize("item1"*100) - with pytest.raises(CosmosResourceNotFoundError): - await container_client.read_item(escaped_key, storage._get_partition_key(escaped_key)) - - escaped_key = storage._sanitize("1230") - item = (await container_client.read_item(escaped_key, storage._get_partition_key(escaped_key))).get("document") - assert ( - MockStoreItemB.from_json_to_store_item(item) - == initial_data["1230"] ) + assert (await cosmos_storage.read(["key"], target_cls=MockStoreItem))[ + "key" + ] == MockStoreItem({"id": "key", "value": "data"}) + assert (await cosmos_storage.read(["key2"], target_cls=MockStoreItem))[ + "key2" + ] == MockStoreItem({"id": "key2", "value": "new_val"}) - @pytest.mark.asyncio - async def test_cosmos_db_storage_flow_existing_container_and_persistence(self): - await self.cosmos_db_storage_flow_existing_container_and_persistence(config) - assert True @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") -class TestCosmosDBStorageWithCompat(dTestCosmosDBStorage): +class TestCosmosDBStorageWithCompat(TestCosmosDBStorage): def get_compat_mode(self): return True - - @pytest.mark.asyncio - async def test_cosmos_db_storage_flow_existing_container_and_persistence(self): - await self.cosmos_db_storage_flow_existing_container_and_persistence(config) - assert True - @pytest.mark.asyncio - async def test_cosmos_db_storage_flow_existing_container_and_persistence_with_compat(self): - await self.cosmos_db_storage_flow_existing_container_and_persistence(config, True) - assert True @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") class TestCosmosDBStorageInit: diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py b/libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py index 8e287e0c..5a2b6f46 100644 --- a/libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py +++ b/libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py @@ -4,6 +4,7 @@ # thank you AI + def test_sanitize_key_simple(): """Test sanitize_key with valid keys that shouldn't change""" assert sanitize_key("validKey123") == "validKey123" @@ -13,6 +14,7 @@ def test_sanitize_key_simple(): assert sanitize_key("with-dash") == "with-dash" assert sanitize_key("with.dot") == "with.dot" + def test_sanitize_key_forbidden_chars(): """Test sanitize_key with forbidden characters""" # Test each forbidden character individually @@ -25,11 +27,13 @@ def test_sanitize_key_forbidden_chars(): assert sanitize_key("key\rvalue") == "key*13value" # carriage return assert sanitize_key("key*value") == "key*42value" # asterisk + def test_sanitize_key_multiple_forbidden_chars(): """Test sanitize_key with multiple forbidden characters""" assert sanitize_key("key/with\\many?bad#chars") == "key*47with*92many*63bad*35chars" assert sanitize_key("a\\b/c?d#e\tf\ng\rh*i") == "a*92b*47c*63d*35e*9f*10g*13h*42i" + def test_sanitize_key_with_long_key_with_forbidden_chars(): """Test sanitize_key with a long key that requires truncation""" long_key = "a?2/!@\t3." * 100 # Create a long key @@ -40,6 +44,7 @@ def test_sanitize_key_with_long_key_with_forbidden_chars(): assert "/" not in sanitized assert "\t" not in sanitized + def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix(): """Test sanitize_key with a long key that requires truncation""" long_key = "a?2/!@\t3." * 100 # Create a long key @@ -50,6 +55,7 @@ def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix(): assert "/" not in sanitized assert "#" not in sanitized + def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix_compat_mode(): """Test sanitize_key with a long key that requires truncation""" long_key = "a?2/!@\t3." * 100 # Create a long key @@ -60,12 +66,14 @@ def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix_compat_mode assert "/" not in sanitized assert "#" not in sanitized + def test_sanitize_key_empty_and_whitespace(): """Test sanitize_key with empty and whitespace strings""" assert sanitize_key("") == "" assert sanitize_key(" ") == " " # spaces are allowed assert sanitize_key(" key ") == " key " # leading/trailing spaces preserved + def test_sanitize_key_with_suffix(): """Test sanitize_key with key_suffix parameter""" assert sanitize_key("key", key_suffix="_suffix") == "key_suffix" @@ -73,25 +81,30 @@ def test_sanitize_key_with_suffix(): assert sanitize_key("key/value", key_suffix="_clean") == "key*47value_clean" assert sanitize_key("", key_suffix="_suffix") == "_suffix" + def test_sanitize_key_suffix_with_truncation(): """Test sanitize_key with suffix that causes truncation""" long_key = "a" * 250 suffix = "_suffix" result = sanitize_key(long_key, key_suffix=suffix, compatibility_mode=True) assert len(result) <= 255 - assert result.endswith(suffix) or len(result) == 255 # Either has suffix or was truncated + assert ( + result.endswith(suffix) or len(result) == 255 + ) # Either has suffix or was truncated + def test_sanitize_key_truncation_compatibility_mode(): """Test sanitize_key with truncation in compatibility mode""" long_key = "a" * 300 result = sanitize_key(long_key, compatibility_mode=True) assert len(result) <= 255 - + # Should contain hash when truncated very_long_key = "b" * 500 result2 = sanitize_key(very_long_key, compatibility_mode=True) assert len(result2) == 255 + def test_sanitize_key_no_truncation(): """Test sanitize_key without compatibility mode (no truncation)""" long_key = "a" * 300 @@ -99,44 +112,49 @@ def test_sanitize_key_no_truncation(): assert result == long_key # Should be unchanged assert len(result) == 300 + def test_truncate_key_short_strings(): """Test truncate_key with strings shorter than 255 characters""" assert truncate_key("short") == "short" assert truncate_key("a" * 254) == "a" * 254 assert truncate_key("a" * 255) == "a" * 255 + def test_truncate_key_long_strings(): """Test truncate_key with strings longer than 255 characters""" long_key = "a" * 300 result = truncate_key(long_key) assert len(result) == 255 - + # Result should end with SHA256 hash expected_hash = hashlib.sha256(long_key.encode("utf-8")).hexdigest() assert result.endswith(expected_hash) - + # First part should be original key truncated expected_prefix_len = 255 - len(expected_hash) assert result.startswith("a" * expected_prefix_len) + def test_truncate_key_compatibility_mode_disabled(): """Test truncate_key with compatibility_mode=False""" long_key = "a" * 300 assert truncate_key(long_key, compatibility_mode=False) == long_key - + very_long_key = "x" * 1000 assert truncate_key(very_long_key, compatibility_mode=False) == very_long_key - + # Should work with any characters when compatibility mode is off special_chars = "key/with\\special?chars#and\ttabs\nand\rmore*" assert truncate_key(special_chars, compatibility_mode=False) == special_chars + def test_truncate_key_exactly_255_chars(): """Test truncate_key with exactly 255 characters""" key_255 = "a" * 255 assert truncate_key(key_255) == key_255 assert len(truncate_key(key_255)) == 255 + def test_truncate_key_256_chars(): """Test truncate_key with 256 characters (one over limit)""" key_256 = "a" * 256 @@ -144,6 +162,7 @@ def test_truncate_key_256_chars(): assert len(result) == 255 assert result != key_256 + def test_truncate_key_hash_consistency(): """Test that truncate_key produces consistent hashes for same input""" long_key = "consistent_test_key_" * 20 # > 255 chars @@ -152,6 +171,7 @@ def test_truncate_key_hash_consistency(): assert result1 == result2 assert len(result1) == 255 + def test_truncate_key_different_inputs_different_outputs(): """Test that different long keys produce different truncated results""" key1 = "a" * 300 @@ -161,38 +181,44 @@ def test_truncate_key_different_inputs_different_outputs(): assert result1 != result2 assert len(result1) == len(result2) == 255 + def test_sanitize_key_integration(): """Integration test combining forbidden chars, suffix, and truncation""" # Key with forbidden chars that will be long after sanitization + suffix base_key = "test/key\\with?many#forbidden\tchars\nand\rmore*" * 10 suffix = "_integration_test" - + result = sanitize_key(base_key, key_suffix=suffix, compatibility_mode=True) - + # Should be sanitized and truncated assert len(result) <= 255 assert "*47" in result or "*92" in result # Contains sanitized chars - + # Test without truncation - result_no_trunc = sanitize_key(base_key, key_suffix=suffix, compatibility_mode=False) - assert "*47" in result_no_trunc or "*92" in result_no_trunc # Contains sanitized chars + result_no_trunc = sanitize_key( + base_key, key_suffix=suffix, compatibility_mode=False + ) + assert ( + "*47" in result_no_trunc or "*92" in result_no_trunc + ) # Contains sanitized chars assert result_no_trunc.endswith(suffix) + def test_edge_cases(): """Test various edge cases""" # Unicode characters (should be preserved) assert sanitize_key("key_ñ_测试") == "key_ñ_测试" - + # Numbers only assert sanitize_key("123456789") == "123456789" - + # Mixed case with special chars result = sanitize_key("MyKey/WithSlash") assert result == "MyKey*47WithSlash" - + # Key that becomes exactly 255 chars after sanitization # Create a key that will be exactly 255 after replacing one char base = "a" * 252 + "/" # 253 chars, "/" becomes "*47" (3 chars) = 255 total result = sanitize_key(base) assert len(result) == 255 - assert result.endswith("*47") \ No newline at end of file + assert result.endswith("*47") diff --git a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/error_handling.py b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/error_handling.py index 5bcd89b4..40a62d75 100644 --- a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/error_handling.py +++ b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/error_handling.py @@ -7,7 +7,7 @@ async def ignore_error(promise: Awaitable, ignore_error_filter: error_filter): """ Ignores errors based on the provided filter function. - + promise: the awaitable to execute ignore_error_filter: a function that takes an Exception and returns True if the error should be ignored, False otherwise. diff --git a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py index edc603dd..e095cbd6 100644 --- a/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py +++ b/libraries/Storage/microsoft-agents-storage/microsoft/agents/storage/storage_test_utils.py @@ -9,6 +9,7 @@ from ._type_aliases import JSON from .memory_storage import MemoryStorage + class MockStoreItem(StoreItem): """Test implementation of StoreItem for testing purposes""" @@ -200,7 +201,7 @@ def changes(self, request): class CRUDStorageTests(StorageTestsCommon): """Tests for Storage implementations that support CRUD operations. - + To use, subclass and implement the `storage` method. """ @@ -449,23 +450,22 @@ async def test_flow(self): storage_alt = await self.storage(existing=True) assert await baseline_storage.equals(storage_alt) + class QuickCRUDStorageTests(CRUDStorageTests): """Reduced set of permutations for quicker tests. Useful for debugging.""" - KEY_LIST = ([ - "\\?/#\t\n\r*", "test.txt" - ]) + KEY_LIST = ["\\?/#\t\n\r*", "test.txt"] READ_KEY_LIST = KEY_LIST + ["nonexistent_key"] STATE_LIST = [ - { key: MockStoreItem({"id": key, "value": f"value{key}"}) for key in KEY_LIST } + {key: MockStoreItem({"id": key, "value": f"value{key}"}) for key in KEY_LIST} ] @pytest.fixture(params=STATE_LIST) def initial_state(self, request): return request.param - + @pytest.fixture(params=KEY_LIST) def key(self, request): return request.param @@ -478,22 +478,29 @@ def keys(self, request): def changes(self, request): changes_obj = {} keys = request.param - changes_obj["new_key"] = MockStoreItemB({"field": "new_value_for_new_key"}, True) + changes_obj["new_key"] = MockStoreItemB( + {"field": "new_value_for_new_key"}, True + ) for i, key in enumerate(keys): if i % 2 == 0: - changes_obj[key] = MockStoreItemB({"data": f"value{key}"}, (i // 2) % 2 == 0) + changes_obj[key] = MockStoreItemB( + {"data": f"value{key}"}, (i // 2) % 2 == 0 + ) else: - changes_obj[key] = MockStoreItem({"id": key, "value": f"new_value_for_{key}"}) + changes_obj[key] = MockStoreItem( + {"id": key, "value": f"new_value_for_{key}"} + ) changes_obj["new_key_2"] = MockStoreItem({"field": "new_value_for_new_key_2"}) return changes_obj + def debug_print(*args): """Print debug information clearly separated in the console.""" - print("\n"*2) + print("\n" * 2) print("--- DEBUG ---") for arg in args: - print("\n"*2) + print("\n" * 2) print(arg) - print("\n"*2) + print("\n" * 2) print("--- ----- ---") - print("\n" *2) \ No newline at end of file + print("\n" * 2) diff --git a/libraries/Storage/microsoft-agents-storage/tests/test_memory_storage.py b/libraries/Storage/microsoft-agents-storage/tests/test_memory_storage.py index f1d15865..de1a5ff5 100644 --- a/libraries/Storage/microsoft-agents-storage/tests/test_memory_storage.py +++ b/libraries/Storage/microsoft-agents-storage/tests/test_memory_storage.py @@ -1,10 +1,11 @@ from microsoft.agents.storage.memory_storage import MemoryStorage from microsoft.agents.storage.storage_test_utils import CRUDStorageTests + class TestMemoryStorage(CRUDStorageTests): async def storage(self, initial_data=None): data = { key: value.store_item_to_json() for key, value in (initial_data or {}).items() } - return MemoryStorage(data) \ No newline at end of file + return MemoryStorage(data) From 989223b8d73e810bb557e5337ff7d7e0ae38288e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Mon, 21 Jul 2025 13:08:08 -0700 Subject: [PATCH 08/13] Added microsoft-agent-cosmos to CI files --- .azdo/ci-pr.yaml | 1 + .github/workflows/python-package.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.azdo/ci-pr.yaml b/.azdo/ci-pr.yaml index 64742e94..2028b9b8 100644 --- a/.azdo/ci-pr.yaml +++ b/.azdo/ci-pr.yaml @@ -58,6 +58,7 @@ steps: python -m pip install ./dist/microsoft_agents_hosting_aiohttp*.whl python -m pip install ./dist/microsoft_agents_storage*.whl python -m pip install ./dist/microsoft_agents_blob*.whl + python -m pip install ./dist/microsoft_agents_cosmos*.whl displayName: 'Install wheels' - script: | diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 5e0b070f..a4e2912e 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -64,6 +64,7 @@ jobs: python -m pip install ./dist/microsoft_agents_hosting_aiohttp*.whl python -m pip install ./dist/microsoft_agents_storage*.whl python -m pip install ./dist/microsoft_agents_blob*.whl + python -m pip install ./dist/microsoft_agents_cosmos*.whl - name: Test with pytest run: | pytest From 0b0fe90ce765f71a8c5b1bcc39fef968cbfe29dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Mon, 21 Jul 2025 13:16:29 -0700 Subject: [PATCH 09/13] Minor fixes to storage tests --- .../tests/test_blob_storage.py | 20 ++++++++----------- .../tests/test_cosmos_db_storage.py | 1 + 2 files changed, 9 insertions(+), 12 deletions(-) diff --git a/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py b/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py index 468c4274..936c799d 100644 --- a/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py +++ b/libraries/Storage/microsoft-agents-blob/tests/test_blob_storage.py @@ -97,12 +97,8 @@ async def test_initialize(self, blob_storage): @pytest.mark.asyncio async def test_external_change_is_visible(self): blob_storage, container_client = await blob_storage_instance() - assert ( - await blob_storage.read(["key"], target_cls=MockStoreItem) - ) == MockStoreItem({"id": "item", "value": "data"}) - assert ( - await blob_storage.read(["key2"], target_cls=MockStoreItem) - ) == MockStoreItem({"id": "another_item", "value": "another_value"}) + assert (await blob_storage.read(["key"], target_cls=MockStoreItem)) == {} + assert (await blob_storage.read(["key2"], target_cls=MockStoreItem)) == {} await container_client.upload_blob( name="key", data=json.dumps({"id": "item", "value": "data"}), overwrite=True ) @@ -111,12 +107,12 @@ async def test_external_change_is_visible(self): data=json.dumps({"id": "another_item", "value": "new_val"}), overwrite=True, ) - assert ( - await blob_storage.read(["key"], target_cls=MockStoreItem) - ) == MockStoreItem({"id": "item", "value": "data"}) - assert ( - await blob_storage.read(["key2"], target_cls=MockStoreItem) - ) == MockStoreItem({"id": "another_item", "value": "new_val"}) + assert (await blob_storage.read(["key"], target_cls=MockStoreItem))[ + "key" + ] == MockStoreItem({"id": "item", "value": "data"}) + assert (await blob_storage.read(["key2"], target_cls=MockStoreItem))[ + "key2" + ] == MockStoreItem({"id": "another_item", "value": "new_val"}) @pytest.mark.asyncio async def test_blob_storage_flow_existing_container_and_persistence(self): diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py index b004a69f..deb3de0d 100644 --- a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py +++ b/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py @@ -102,6 +102,7 @@ async def cosmos_db_storage(): @pytest.mark.asyncio @pytest.mark.parametrize("test_require_compat", [True, False]) +@pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") async def test_cosmos_db_storage_flow_existing_container_and_persistence( test_require_compat, ): From f943f8c8361486c29afee1589f5df9781127727b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Wed, 23 Jul 2025 15:04:15 -0700 Subject: [PATCH 10/13] Moved files around --- .../agents/storage}/cosmos/__init__.py | 0 .../storage}/cosmos/cosmos_db_storage.py | 43 +--- .../cosmos/cosmos_db_storage_config.py | 188 +++++++++--------- .../agents/storage}/cosmos/key_ops.py | 90 ++++----- .../pyproject.toml | 6 +- .../tests/test_cosmos_db_config.py | 2 +- .../tests/test_cosmos_db_storage.py | 7 +- .../tests/test_key_ops.py | 0 8 files changed, 153 insertions(+), 183 deletions(-) rename libraries/{Storage/microsoft-agents-cosmos/microsoft/agents => microsoft-agents-storage-cosmos/microsoft/agents/storage}/cosmos/__init__.py (100%) rename libraries/{Storage/microsoft-agents-cosmos/microsoft/agents => microsoft-agents-storage-cosmos/microsoft/agents/storage}/cosmos/cosmos_db_storage.py (82%) rename libraries/{Storage/microsoft-agents-cosmos/microsoft/agents => microsoft-agents-storage-cosmos/microsoft/agents/storage}/cosmos/cosmos_db_storage_config.py (97%) rename libraries/{Storage/microsoft-agents-cosmos/microsoft/agents => microsoft-agents-storage-cosmos/microsoft/agents/storage}/cosmos/key_ops.py (94%) rename libraries/{Storage/microsoft-agents-cosmos => microsoft-agents-storage-cosmos}/pyproject.toml (77%) rename libraries/{Storage/microsoft-agents-cosmos => microsoft-agents-storage-cosmos}/tests/test_cosmos_db_config.py (99%) rename libraries/{Storage/microsoft-agents-cosmos => microsoft-agents-storage-cosmos}/tests/test_cosmos_db_storage.py (98%) rename libraries/{Storage/microsoft-agents-cosmos => microsoft-agents-storage-cosmos}/tests/test_key_ops.py (100%) diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/__init__.py b/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/__init__.py similarity index 100% rename from libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/__init__.py rename to libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/__init__.py diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py b/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage.py similarity index 82% rename from libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py rename to libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage.py index aa6a2bd2..e276ecce 100644 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage.py +++ b/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage.py @@ -19,7 +19,7 @@ from microsoft.agents.storage import AsyncStorageBase, StoreItem from microsoft.agents.storage._type_aliases import JSON -from microsoft.agents.storage.error_handling import ignore_error, is_status_code_error +from microsoft.agents.storage.error_handling import ignore_error from .cosmos_db_storage_config import CosmosDBStorageConfig from .key_ops import sanitize_key @@ -84,38 +84,6 @@ def _sanitize(self, key: str) -> str: key, self._config.key_suffix, self._config.compatibility_mode ) - async def _get_item_safe( - self, raw_key: str, sanitized_key: str - ) -> Union[JSON, None]: - """Get item, ensuring no collision occurs. - - :param raw_key: The original key provided by the user. - :param sanitized_key: The sanitized key used for CosmosDB operations. - :return: The item if found, None otherwise. - - It's highly highly highly (highly) unlikely that a sha256 collision will ever occur. - But, just in case. - """ - - read_item_response: CosmosDict = await ignore_error( - self._container.read_item( - sanitized_key, self._get_partition_key(sanitized_key) - ), - cosmos_resource_not_found, - ) - - if ( - read_item_response is not None - and read_item_response.get("realId") is not None - and read_item_response.get("realId") != raw_key - ): - # This should never happen, but just in case... - # in fact, if this does happen, then we should probably document this - # and publish the first ever sha256 collision to the world. - raise Exception("CosmosDBStorage: Key mismatch on get_item_raw.") - - return read_item_response - async def _read_item( self, key: str, *, target_cls: StoreItemT = None, **kwargs ) -> tuple[Union[str, None], Union[StoreItemT, None]]: @@ -124,7 +92,12 @@ async def _read_item( raise ValueError("CosmosDBStorage: Key cannot be empty.") escaped_key: str = self._sanitize(key) - read_item_response: CosmosDict = await self._get_item_safe(key, escaped_key) + read_item_response: CosmosDict = await ignore_error( + self._container.read_item( + escaped_key, self._get_partition_key(escaped_key) + ), + cosmos_resource_not_found, + ) if read_item_response is None: return None, None @@ -136,7 +109,6 @@ async def _write_item(self, key: str, item: StoreItem) -> None: raise ValueError("CosmosDBStorage: Key cannot be empty.") escaped_key: str = self._sanitize(key) - await self._get_item_safe(key, escaped_key) # ensure no collision doc = { "id": escaped_key, @@ -150,7 +122,6 @@ async def _delete_item(self, key: str) -> None: raise ValueError("CosmosDBStorage: Key cannot be empty.") escaped_key: str = self._sanitize(key) - await self._get_item_safe(key, escaped_key) # ensure no collision await ignore_error( self._container.delete_item( diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py b/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage_config.py similarity index 97% rename from libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py rename to libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage_config.py index 1ebddb5f..e70ec138 100644 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/cosmos_db_storage_config.py +++ b/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage_config.py @@ -1,94 +1,94 @@ -import json -from typing import Union - -from azure.core.credentials import TokenCredential - -from .key_ops import sanitize_key - - -class CosmosDBStorageConfig: - """The class for partitioned CosmosDB configuration for the Azure Bot Framework.""" - - def __init__( - self, - cosmos_db_endpoint: str = "", - auth_key: str = "", - database_id: str = "", - container_id: str = "", - cosmos_client_options: dict = None, - container_throughput: int = 0, - key_suffix: str = "", - compatibility_mode: bool = False, - url: str = "", - credential: Union[TokenCredential, None] = None, - **kwargs, - ): - """Create the Config object. - - :param cosmos_db_endpoint: The CosmosDB endpoint. - :param auth_key: The authentication key for Cosmos DB. - :param database_id: The database identifier for Cosmos DB instance. - :param container_id: The container identifier. - :param cosmos_client_options: The options for the CosmosClient. Currently only supports connection_policy and - consistency_level - :param container_throughput: The throughput set when creating the Container. Defaults to 400. - :param key_suffix: The suffix to be added to every key. The keySuffix must contain only valid ComosDb - key characters. (e.g. not: '\\', '?', '/', '#', '*') - :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb - max key length of 255. - :param url: The URL to the CosmosDB resource. - :param credential: The TokenCredential to use for authentication. - :return CosmosDBConfig: - """ - config_file: str = kwargs.get("filename", "") - if config_file: - kwargs = json.load(open(config_file)) - self.cosmos_db_endpoint: str = cosmos_db_endpoint or kwargs.get( - "cosmos_db_endpoint", "" - ) - self.auth_key: str = auth_key or kwargs.get("auth_key", "") - self.database_id: str = database_id or kwargs.get("database_id", "") - self.container_id: str = container_id or kwargs.get("container_id", "") - self.cosmos_client_options: dict = cosmos_client_options or kwargs.get( - "cosmos_client_options", {} - ) - self.container_throughput: int = container_throughput or kwargs.get( - "container_throughput", 400 - ) - self.key_suffix: str = key_suffix or kwargs.get("key_suffix", "") - self.compatibility_mode: bool = compatibility_mode or kwargs.get( - "compatibility_mode", False - ) - self.url = url or kwargs.get("url", "") - self.credential: Union[TokenCredential, None] = credential - - @staticmethod - def validate_cosmos_db_config(config: "CosmosDBStorageConfig") -> None: - """Validate the CosmosDBConfig object. - - This is used prior to the creation of the CosmosDBStorage object.""" - if not config: - raise ValueError("CosmosDBStorage: CosmosDBConfig is required.") - if not config.cosmos_db_endpoint: - raise ValueError("CosmosDBStorage: cosmos_db_endpoint is required.") - if not config.auth_key: - raise ValueError("CosmosDBStorage: auth_key is required.") - if not config.database_id: - raise ValueError("CosmosDBStorage: database_id is required.") - if not config.container_id: - raise ValueError("CosmosDBStorage: container_id is required.") - - CosmosDBStorageConfig._validate_suffix(config) - - @staticmethod - def _validate_suffix(config: "CosmosDBStorageConfig") -> None: - if config.key_suffix: - if config.compatibility_mode: - raise ValueError( - "compatibilityMode cannot be true while using a keySuffix." - ) - suffix_escaped: str = sanitize_key(config.key_suffix) - if suffix_escaped != config.key_suffix: - raise ValueError( - f"Cannot use invalid Row Key characters: {config.key_suffix} in keySuffix." - ) +import json +from typing import Union + +from azure.core.credentials import TokenCredential + +from .key_ops import sanitize_key + + +class CosmosDBStorageConfig: + """The class for partitioned CosmosDB configuration for the Azure Bot Framework.""" + + def __init__( + self, + cosmos_db_endpoint: str = "", + auth_key: str = "", + database_id: str = "", + container_id: str = "", + cosmos_client_options: dict = None, + container_throughput: int = 0, + key_suffix: str = "", + compatibility_mode: bool = False, + url: str = "", + credential: Union[TokenCredential, None] = None, + **kwargs, + ): + """Create the Config object. + + :param cosmos_db_endpoint: The CosmosDB endpoint. + :param auth_key: The authentication key for Cosmos DB. + :param database_id: The database identifier for Cosmos DB instance. + :param container_id: The container identifier. + :param cosmos_client_options: The options for the CosmosClient. Currently only supports connection_policy and + consistency_level + :param container_throughput: The throughput set when creating the Container. Defaults to 400. + :param key_suffix: The suffix to be added to every key. The keySuffix must contain only valid ComosDb + key characters. (e.g. not: '\\', '?', '/', '#', '*') + :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb + max key length of 255. + :param url: The URL to the CosmosDB resource. + :param credential: The TokenCredential to use for authentication. + :return CosmosDBConfig: + """ + config_file: str = kwargs.get("filename", "") + if config_file: + kwargs = json.load(open(config_file)) + self.cosmos_db_endpoint: str = cosmos_db_endpoint or kwargs.get( + "cosmos_db_endpoint", "" + ) + self.auth_key: str = auth_key or kwargs.get("auth_key", "") + self.database_id: str = database_id or kwargs.get("database_id", "") + self.container_id: str = container_id or kwargs.get("container_id", "") + self.cosmos_client_options: dict = cosmos_client_options or kwargs.get( + "cosmos_client_options", {} + ) + self.container_throughput: int = container_throughput or kwargs.get( + "container_throughput", 400 + ) + self.key_suffix: str = key_suffix or kwargs.get("key_suffix", "") + self.compatibility_mode: bool = compatibility_mode or kwargs.get( + "compatibility_mode", False + ) + self.url = url or kwargs.get("url", "") + self.credential: Union[TokenCredential, None] = credential + + @staticmethod + def validate_cosmos_db_config(config: "CosmosDBStorageConfig") -> None: + """Validate the CosmosDBConfig object. + + This is used prior to the creation of the CosmosDBStorage object.""" + if not config: + raise ValueError("CosmosDBStorage: CosmosDBConfig is required.") + if not config.cosmos_db_endpoint: + raise ValueError("CosmosDBStorage: cosmos_db_endpoint is required.") + if not config.auth_key: + raise ValueError("CosmosDBStorage: auth_key is required.") + if not config.database_id: + raise ValueError("CosmosDBStorage: database_id is required.") + if not config.container_id: + raise ValueError("CosmosDBStorage: container_id is required.") + + CosmosDBStorageConfig._validate_suffix(config) + + @staticmethod + def _validate_suffix(config: "CosmosDBStorageConfig") -> None: + if config.key_suffix: + if config.compatibility_mode: + raise ValueError( + "compatibilityMode cannot be true while using a keySuffix." + ) + suffix_escaped: str = sanitize_key(config.key_suffix) + if suffix_escaped != config.key_suffix: + raise ValueError( + f"Cannot use invalid Row Key characters: {config.key_suffix} in keySuffix." + ) diff --git a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py b/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/key_ops.py similarity index 94% rename from libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py rename to libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/key_ops.py index 1d6941f9..8e91a598 100644 --- a/libraries/Storage/microsoft-agents-cosmos/microsoft/agents/cosmos/key_ops.py +++ b/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/key_ops.py @@ -1,45 +1,45 @@ -from hashlib import sha256 - - -def sanitize_key( - key: str, key_suffix: str = "", compatibility_mode: bool = True -) -> str: - """Return the sanitized key. - - Replace characters that are not allowed in keys in Cosmos. - - :param key: The provided key to be escaped. - :param key_suffix: The string to add a the end of all RowKeys. - :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb - max key length of 255. This behavior can be overridden by setting - cosmosdb_config.compatibility_mode to False. - :return str: - """ - # forbidden characters - bad_chars: list[str] = ["\\", "?", "/", "#", "\t", "\n", "\r", "*"] - - # replace those with with '*' and the - # Unicode code point of the character and return the new string - key = "".join(map(lambda x: "*" + str(ord(x)) if x in bad_chars else x, key)) - return truncate_key(f"{key}{key_suffix}", compatibility_mode) - - -def truncate_key(key: str, compatibility_mode: bool = True) -> str: - """ - Truncate the key to 255 characters if compatibility_mode is True. If the key is longer than 255 characters, - it will be truncated and a SHA-256 hash of the original key will be appended to minimize collisions. - """ - max_key_len: int = 255 - - if not compatibility_mode: - return key - - if len(key) > max_key_len: - # for now - # https://stackoverflow.com/questions/4014090/is-it-safe-to-ignore-the-possibility-of-sha-collisions-in-practice - aux_hash = sha256(key.encode("utf-8")) - aux_hex = aux_hash.hexdigest() - - key = key[0 : max_key_len - len(aux_hex)] + aux_hex - - return key +from hashlib import sha256 + + +def sanitize_key( + key: str, key_suffix: str = "", compatibility_mode: bool = True +) -> str: + """Return the sanitized key. + + Replace characters that are not allowed in keys in Cosmos. + + :param key: The provided key to be escaped. + :param key_suffix: The string to add a the end of all RowKeys. + :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb + max key length of 255. This behavior can be overridden by setting + cosmosdb_config.compatibility_mode to False. + :return str: + """ + # forbidden characters + bad_chars: list[str] = ["\\", "?", "/", "#", "\t", "\n", "\r", "*"] + + # replace those with with '*' and the + # Unicode code point of the character and return the new string + key = "".join(map(lambda x: "*" + str(ord(x)) if x in bad_chars else x, key)) + return truncate_key(f"{key}{key_suffix}", compatibility_mode) + + +def truncate_key(key: str, compatibility_mode: bool = True) -> str: + """ + Truncate the key to 255 characters if compatibility_mode is True. If the key is longer than 255 characters, + it will be truncated and a SHA-256 hash of the original key will be appended to minimize collisions. + """ + max_key_len: int = 255 + + if not compatibility_mode: + return key + + if len(key) > max_key_len: + # for now (and the foreseeable future), SHA-256 collisions are pretty infentesimally rare: + # https://stackoverflow.com/questions/4014090/is-it-safe-to-ignore-the-possibility-of-sha-collisions-in-practice + aux_hash = sha256(key.encode("utf-8")) + aux_hex = aux_hash.hexdigest() + + key = key[0 : max_key_len - len(aux_hex)] + aux_hex + + return key diff --git a/libraries/Storage/microsoft-agents-cosmos/pyproject.toml b/libraries/microsoft-agents-storage-cosmos/pyproject.toml similarity index 77% rename from libraries/Storage/microsoft-agents-cosmos/pyproject.toml rename to libraries/microsoft-agents-storage-cosmos/pyproject.toml index 152df9f4..611d6e46 100644 --- a/libraries/Storage/microsoft-agents-cosmos/pyproject.toml +++ b/libraries/microsoft-agents-storage-cosmos/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools"] build-backend = "setuptools.build_meta" [project] -name = "microsoft-agents-cosmos" +name = "microsoft-agents-storage-cosmos" version = "0.0.0a1" description = "A Cosmos DB storage library for Microsoft Agents" authors = [{name = "Microsoft Corporation"}] @@ -14,10 +14,10 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ - "microsoft.agents.storage", + "microsoft-agents-hosting-core", "azure-core", "azure-cosmos", ] [project.urls] -"Homepage" = "https://github.com/microsoft/microsoft-agents-protocol" +"Homepage" = "https://github.com/microsoft/Agents" diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py b/libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_config.py similarity index 99% rename from libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py rename to libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_config.py index 81f79202..02e47b30 100644 --- a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_config.py +++ b/libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_config.py @@ -1,7 +1,7 @@ import json import pytest -from microsoft.agents.cosmos import CosmosDBStorageConfig +from microsoft.agents.storage.cosmos import CosmosDBStorageConfig # thank you AI, again diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py b/libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_storage.py similarity index 98% rename from libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py rename to libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_storage.py index deb3de0d..71657506 100644 --- a/libraries/Storage/microsoft-agents-cosmos/tests/test_cosmos_db_storage.py +++ b/libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_storage.py @@ -1,8 +1,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -import json, gc -from io import BytesIO +import gc import pytest import pytest_asyncio @@ -11,8 +10,8 @@ from azure.cosmos.aio import CosmosClient from azure.cosmos.exceptions import CosmosResourceNotFoundError -from microsoft.agents.cosmos import CosmosDBStorage, CosmosDBStorageConfig -from microsoft.agents.cosmos.key_ops import sanitize_key +from microsoft.agents.storage.cosmos import CosmosDBStorage, CosmosDBStorageConfig +from microsoft.agents.storage.cosmos.key_ops import sanitize_key from microsoft.agents.storage.storage_test_utils import ( QuickCRUDStorageTests, diff --git a/libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py b/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py similarity index 100% rename from libraries/Storage/microsoft-agents-cosmos/tests/test_key_ops.py rename to libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py From 7aaf176b0a04ffe6bc6ceaeaebea161373ac22e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Wed, 23 Jul 2025 15:19:37 -0700 Subject: [PATCH 11/13] Fixed imports and formatted with black --- .../tests/test_memory_storage.py | 1 + .../agents/storage/blob/blob_storage.py | 2 +- .../storage/cosmos/cosmos_db_storage.py | 6 ++--- .../tests/test_cosmos_db_storage.py | 22 +------------------ .../tests/test_key_ops.py | 2 +- 5 files changed, 7 insertions(+), 26 deletions(-) diff --git a/libraries/microsoft-agents-hosting-core/tests/test_memory_storage.py b/libraries/microsoft-agents-hosting-core/tests/test_memory_storage.py index 70676f34..421af72a 100644 --- a/libraries/microsoft-agents-hosting-core/tests/test_memory_storage.py +++ b/libraries/microsoft-agents-hosting-core/tests/test_memory_storage.py @@ -1,6 +1,7 @@ from microsoft.agents.hosting.core.storage.memory_storage import MemoryStorage from microsoft.agents.hosting.core.storage.storage_test_utils import CRUDStorageTests + class TestMemoryStorage(CRUDStorageTests): async def storage(self, initial_data=None): data = { diff --git a/libraries/microsoft-agents-storage-blob/microsoft/agents/storage/blob/blob_storage.py b/libraries/microsoft-agents-storage-blob/microsoft/agents/storage/blob/blob_storage.py index 818cc3ca..02350a29 100644 --- a/libraries/microsoft-agents-storage-blob/microsoft/agents/storage/blob/blob_storage.py +++ b/libraries/microsoft-agents-storage-blob/microsoft/agents/storage/blob/blob_storage.py @@ -97,4 +97,4 @@ async def _write_item(self, key: str, item: StoreItem) -> None: async def _delete_item(self, key: str) -> None: await ignore_error( self._container_client.delete_blob(blob=key), is_status_code_error(404) - ) \ No newline at end of file + ) diff --git a/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage.py b/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage.py index e276ecce..ea589c62 100644 --- a/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage.py +++ b/libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/cosmos_db_storage.py @@ -17,9 +17,9 @@ import azure.cosmos.exceptions as cosmos_exceptions from azure.cosmos.partition_key import NonePartitionKeyValue -from microsoft.agents.storage import AsyncStorageBase, StoreItem -from microsoft.agents.storage._type_aliases import JSON -from microsoft.agents.storage.error_handling import ignore_error +from microsoft.agents.hosting.core.storage import AsyncStorageBase, StoreItem +from microsoft.agents.hosting.core.storage._type_aliases import JSON +from microsoft.agents.hosting.core.storage.error_handling import ignore_error from .cosmos_db_storage_config import CosmosDBStorageConfig from .key_ops import sanitize_key diff --git a/libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_storage.py b/libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_storage.py index 71657506..cd528db1 100644 --- a/libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_storage.py +++ b/libraries/microsoft-agents-storage-cosmos/tests/test_cosmos_db_storage.py @@ -13,7 +13,7 @@ from microsoft.agents.storage.cosmos import CosmosDBStorage, CosmosDBStorageConfig from microsoft.agents.storage.cosmos.key_ops import sanitize_key -from microsoft.agents.storage.storage_test_utils import ( +from microsoft.agents.hosting.core.storage.storage_test_utils import ( QuickCRUDStorageTests, MockStoreItem, MockStoreItemB, @@ -198,26 +198,6 @@ async def test_initialize(self, cosmos_db_storage): await cosmos_db_storage.read(["some_Key"], target_cls=MockStoreItem) ) == {"some_Key": MockStoreItem({"id": "123", "data": "value"})} - @pytest.mark.asyncio - async def test_collision_detection(self): - cosmos_storage, container_client = await cosmos_db_storage_instance() - await container_client.upsert_item( - { - "id": "key", - "realId": "fake_key", - "document": {"id": "key", "value": "data"}, - "partitionKey": "", - } - ) - with pytest.raises(Exception): - await cosmos_storage.read(["key"], target_cls=MockStoreItem) - with pytest.raises(Exception): - await cosmos_storage.write( - {"key": MockStoreItem({"id": "item", "value": "data"})} - ) - with pytest.raises(Exception): - await cosmos_storage.delete(["key"]) - @pytest.mark.asyncio async def test_external_change_is_visible(self): cosmos_storage, container_client = await cosmos_db_storage_instance() diff --git a/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py b/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py index 5a2b6f46..8dff3dba 100644 --- a/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py +++ b/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py @@ -1,6 +1,6 @@ import hashlib -from microsoft.agents.cosmos.key_ops import truncate_key, sanitize_key +from microsoft.agents.storage.cosmos.key_ops import truncate_key, sanitize_key # thank you AI From ec482e75cffbe6054759b6d0dfd2e27e3ae426f9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Thu, 24 Jul 2025 08:59:04 -0700 Subject: [PATCH 12/13] Made key operation tests more data driven --- .../tests/test_key_ops.py | 219 ++++++++++-------- 1 file changed, 119 insertions(+), 100 deletions(-) diff --git a/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py b/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py index 8dff3dba..d0228055 100644 --- a/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py +++ b/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py @@ -1,41 +1,55 @@ import hashlib - +import pytest from microsoft.agents.storage.cosmos.key_ops import truncate_key, sanitize_key # thank you AI -def test_sanitize_key_simple(): - """Test sanitize_key with valid keys that shouldn't change""" - assert sanitize_key("validKey123") == "validKey123" - assert sanitize_key("simple") == "simple" - assert sanitize_key("CamelCase") == "CamelCase" - assert sanitize_key("under_score") == "under_score" - assert sanitize_key("with-dash") == "with-dash" - assert sanitize_key("with.dot") == "with.dot" - - -def test_sanitize_key_forbidden_chars(): - """Test sanitize_key with forbidden characters""" - # Test each forbidden character individually - assert sanitize_key("key\\value") == "key*92value" # backslash - assert sanitize_key("key?value") == "key*63value" # question mark - assert sanitize_key("key/value") == "key*47value" # forward slash - assert sanitize_key("key#value") == "key*35value" # hash - assert sanitize_key("key\tvalue") == "key*9value" # tab - assert sanitize_key("key\nvalue") == "key*10value" # newline - assert sanitize_key("key\rvalue") == "key*13value" # carriage return - assert sanitize_key("key*value") == "key*42value" # asterisk - - -def test_sanitize_key_multiple_forbidden_chars(): - """Test sanitize_key with multiple forbidden characters""" - assert sanitize_key("key/with\\many?bad#chars") == "key*47with*92many*63bad*35chars" - assert sanitize_key("a\\b/c?d#e\tf\ng\rh*i") == "a*92b*47c*63d*35e*9f*10g*13h*42i" +@pytest.mark.parametrize( + "input_key,expected", + [ + ("validKey123", "validKey123"), + ("simple", "simple"), + ("CamelCase", "CamelCase"), + ("under_score", "under_score"), + ("with-dash", "with-dash"), + ("with.dot", "with.dot"), + ], +) +def test_sanitize_key_simple(input_key, expected): + assert sanitize_key(input_key) == expected + + +@pytest.mark.parametrize( + "input_key,expected,description", + [ + ("key\\value", "key*92value"), + ("key?value", "key*63value"), + ("key/value", "key*47value"), + ("key#value", "key*35value"), + ("key\tvalue", "key*9value"), + ("key\nvalue", "key*10value"), + ("key\rvalue", "key*13value"), + ("key*value", "key*42value"), + ], +) +def test_sanitize_key_forbidden_chars(input_key, expected): + assert sanitize_key(input_key) == expected + + +@pytest.mark.parametrize( + "input_key,expected", + [ + ("key/with\\many?bad#chars", "key*47with*92many*63bad*35chars"), + ("a\\b/c?d#e\tf\ng\rh*i", "a*92b*47c*63d*35e*9f*10g*13h*42i"), + ("key/with\\many?bad#chars", "key*47with*92many*63bad*35chars"), + ], +) +def test_sanitize_key_multiple_forbidden_chars(input_key, expected): + assert sanitize_key(input_key) == expected def test_sanitize_key_with_long_key_with_forbidden_chars(): - """Test sanitize_key with a long key that requires truncation""" long_key = "a?2/!@\t3." * 100 # Create a long key sanitized = sanitize_key(long_key) assert len(sanitized) <= 255 # Should be truncated @@ -46,7 +60,6 @@ def test_sanitize_key_with_long_key_with_forbidden_chars(): def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix(): - """Test sanitize_key with a long key that requires truncation""" long_key = "a?2/!@\t3." * 100 # Create a long key sanitized = sanitize_key(long_key, key_suffix="_suff?#*") assert len(sanitized) <= 255 # Should be truncated @@ -57,7 +70,6 @@ def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix(): def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix_compat_mode(): - """Test sanitize_key with a long key that requires truncation""" long_key = "a?2/!@\t3." * 100 # Create a long key sanitized = sanitize_key(long_key, key_suffix="_suff?#*", compatibility_mode=True) assert len(sanitized) <= 255 # Should be truncated @@ -68,22 +80,23 @@ def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix_compat_mode def test_sanitize_key_empty_and_whitespace(): - """Test sanitize_key with empty and whitespace strings""" - assert sanitize_key("") == "" - assert sanitize_key(" ") == " " # spaces are allowed - assert sanitize_key(" key ") == " key " # leading/trailing spaces preserved + assert sanitize_key(input_key) == expected -def test_sanitize_key_with_suffix(): - """Test sanitize_key with key_suffix parameter""" - assert sanitize_key("key", key_suffix="_suffix") == "key_suffix" - assert sanitize_key("test", key_suffix="123") == "test123" - assert sanitize_key("key/value", key_suffix="_clean") == "key*47value_clean" - assert sanitize_key("", key_suffix="_suffix") == "_suffix" +@pytest.mark.parametrize( + "input_key,suffix,expected", + [ + ("key", "_suffix", "key_suffix"), + ("test", "123", "test123"), + ("key/value", "_clean", "key*47value_clean"), + ("", "_suffix", "_suffix"), + ], +) +def test_sanitize_key_with_suffix(input_key, suffix, expected): + assert sanitize_key(input_key, key_suffix=suffix) == expected def test_sanitize_key_suffix_with_truncation(): - """Test sanitize_key with suffix that causes truncation""" long_key = "a" * 250 suffix = "_suffix" result = sanitize_key(long_key, key_suffix=suffix, compatibility_mode=True) @@ -94,7 +107,6 @@ def test_sanitize_key_suffix_with_truncation(): def test_sanitize_key_truncation_compatibility_mode(): - """Test sanitize_key with truncation in compatibility mode""" long_key = "a" * 300 result = sanitize_key(long_key, compatibility_mode=True) assert len(result) <= 255 @@ -106,22 +118,25 @@ def test_sanitize_key_truncation_compatibility_mode(): def test_sanitize_key_no_truncation(): - """Test sanitize_key without compatibility mode (no truncation)""" long_key = "a" * 300 result = sanitize_key(long_key, compatibility_mode=False) assert result == long_key # Should be unchanged assert len(result) == 300 -def test_truncate_key_short_strings(): - """Test truncate_key with strings shorter than 255 characters""" - assert truncate_key("short") == "short" - assert truncate_key("a" * 254) == "a" * 254 - assert truncate_key("a" * 255) == "a" * 255 +@pytest.mark.parametrize( + "input_key,expected", + [ + ("short", "short"), + ("a" * 254, "a" * 254), + ("a" * 255, "a" * 255), + ], +) +def test_truncate_key_short_strings(input_key, expected): + assert truncate_key(input_key) == expected def test_truncate_key_long_strings(): - """Test truncate_key with strings longer than 255 characters""" long_key = "a" * 300 result = truncate_key(long_key) assert len(result) == 255 @@ -135,36 +150,44 @@ def test_truncate_key_long_strings(): assert result.startswith("a" * expected_prefix_len) -def test_truncate_key_compatibility_mode_disabled(): - """Test truncate_key with compatibility_mode=False""" - long_key = "a" * 300 - assert truncate_key(long_key, compatibility_mode=False) == long_key - - very_long_key = "x" * 1000 - assert truncate_key(very_long_key, compatibility_mode=False) == very_long_key - - # Should work with any characters when compatibility mode is off - special_chars = "key/with\\special?chars#and\ttabs\nand\rmore*" - assert truncate_key(special_chars, compatibility_mode=False) == special_chars - - -def test_truncate_key_exactly_255_chars(): - """Test truncate_key with exactly 255 characters""" - key_255 = "a" * 255 - assert truncate_key(key_255) == key_255 - assert len(truncate_key(key_255)) == 255 - - -def test_truncate_key_256_chars(): - """Test truncate_key with 256 characters (one over limit)""" - key_256 = "a" * 256 - result = truncate_key(key_256) - assert len(result) == 255 - assert result != key_256 +@pytest.mark.parametrize( + "input_key,compatibility_mode,expected_unchanged", + [ + ("a" * 300, False, True), # Should be unchanged + ("x" * 1000, False, True), # Should be unchanged + ( + "key/with\\special?chars#and\ttabs\nand\rmore*", + False, + True, + ), # Should be unchanged + ], +) +def test_truncate_key_compatibility_mode_disabled( + input_key, compatibility_mode, expected_unchanged +): + result = truncate_key(input_key, compatibility_mode=compatibility_mode) + if expected_unchanged: + assert result == input_key + + +@pytest.mark.parametrize( + "input_key,expected_length", + [ + ("a" * 255, 255), + ("a" * 256, 255), + ], +) +def test_truncate_key_exact_and_over_limit(input_key, expected_length): + result = truncate_key(input_key) + assert len(result) == expected_length + + if len(input_key) == 255: + assert result == input_key + else: + assert result != input_key def test_truncate_key_hash_consistency(): - """Test that truncate_key produces consistent hashes for same input""" long_key = "consistent_test_key_" * 20 # > 255 chars result1 = truncate_key(long_key) result2 = truncate_key(long_key) @@ -172,10 +195,14 @@ def test_truncate_key_hash_consistency(): assert len(result1) == 255 -def test_truncate_key_different_inputs_different_outputs(): - """Test that different long keys produce different truncated results""" - key1 = "a" * 300 - key2 = "b" * 300 +@pytest.mark.parametrize( + "key1,key2", + [ + ("a" * 300, "b" * 300), + ("consistent_test_key_" * 20, "different_test_key_" * 20), + ], +) +def test_truncate_key_different_inputs_different_outputs(key1, key2): result1 = truncate_key(key1) result2 = truncate_key(key2) assert result1 != result2 @@ -183,7 +210,6 @@ def test_truncate_key_different_inputs_different_outputs(): def test_sanitize_key_integration(): - """Integration test combining forbidden chars, suffix, and truncation""" # Key with forbidden chars that will be long after sanitization + suffix base_key = "test/key\\with?many#forbidden\tchars\nand\rmore*" * 10 suffix = "_integration_test" @@ -204,21 +230,14 @@ def test_sanitize_key_integration(): assert result_no_trunc.endswith(suffix) -def test_edge_cases(): - """Test various edge cases""" - # Unicode characters (should be preserved) - assert sanitize_key("key_ñ_测试") == "key_ñ_测试" - - # Numbers only - assert sanitize_key("123456789") == "123456789" - - # Mixed case with special chars - result = sanitize_key("MyKey/WithSlash") - assert result == "MyKey*47WithSlash" - - # Key that becomes exactly 255 chars after sanitization - # Create a key that will be exactly 255 after replacing one char - base = "a" * 252 + "/" # 253 chars, "/" becomes "*47" (3 chars) = 255 total - result = sanitize_key(base) - assert len(result) == 255 - assert result.endswith("*47") +@pytest.mark.parametrize( + "input_key,expected,description", + [ + ("key_ñ_测试", "key_ñ_测试"), + ("123456789", "123456789"), + ("MyKey/WithSlash", "MyKey*47WithSlash"), + ], +) +def test_edge_cases(input_key, expected): + result = sanitize_key(input_key) + assert result == expected From ab99fe8362988dea8793901c8c5e053c4bff7af2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Brand=C3=A3o?= Date: Thu, 24 Jul 2025 09:05:17 -0700 Subject: [PATCH 13/13] Fixed minor issues in unit tests --- .../tests/test_key_ops.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py b/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py index d0228055..e554370f 100644 --- a/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py +++ b/libraries/microsoft-agents-storage-cosmos/tests/test_key_ops.py @@ -21,7 +21,7 @@ def test_sanitize_key_simple(input_key, expected): @pytest.mark.parametrize( - "input_key,expected,description", + "input_key,expected", [ ("key\\value", "key*92value"), ("key?value", "key*63value"), @@ -79,7 +79,14 @@ def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix_compat_mode assert "#" not in sanitized -def test_sanitize_key_empty_and_whitespace(): +@pytest.mark.parametrize( + "input_key,expected", + [ + ("", ""), + (" ", " "), + ], +) +def test_sanitize_key_empty_and_whitespace(input_key, expected): assert sanitize_key(input_key) == expected @@ -231,7 +238,7 @@ def test_sanitize_key_integration(): @pytest.mark.parametrize( - "input_key,expected,description", + "input_key,expected", [ ("key_ñ_测试", "key_ñ_测试"), ("123456789", "123456789"),