Adding Cosmos DB support for storage (#68)

* Adding Cosmos DB storage

* Removed unnecessary file

* Added quick storage testing functionality for debugging and improved documentation

* Simplified storage test setup and added unit test

* Fixed legacy container support and added relevant tests

* Revised AsyncStorageBase export/imports

* Fixed final unit tests and refined documentation

* Added microsoft-agent-cosmos to CI files

* Minor fixes to storage tests

* Moved files around

* Fixed imports and formatted with black

* Made key operation tests more data driven

* Fixed minor issues in unit tests

Author: rodrigobr-msft

.azdo/ci-pr.yaml

@@ -51,8 +51,9 @@ steps:
     python -m pip install ./dist/microsoft_agents_hosting_aiohttp*.whl
     python -m pip install ./dist/microsoft_agents_hosting_teams*.whl
     python -m pip install ./dist/microsoft_agents_storage_blob*.whl
+    python -m pip install ./dist/microsoft_agents_storage_cosmos*.whl
   displayName: 'Install wheels'
 
 - script: |
     pytest
-  displayName: 'Test with pytest'
+  displayName: 'Test with pytest'

.github/workflows/python-package.yml

@@ -57,6 +57,7 @@ jobs:
         python -m pip install ./dist/microsoft_agents_hosting_aiohttp*.whl
         python -m pip install ./dist/microsoft_agents_hosting_teams*.whl
         python -m pip install ./dist/microsoft_agents_storage_blob*.whl
+        python -m pip install ./dist/microsoft_agents_storage_cosmos*.whl
     - name: Test with pytest
       run: |
         pytest

libraries/microsoft-agents-hosting-core/microsoft/agents/hosting/core/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"]

libraries/microsoft-agents-hosting-core/microsoft/agents/hosting/core/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:

libraries/microsoft-agents-hosting-core/microsoft/agents/hosting/core/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"""

libraries/microsoft-agents-hosting-core/microsoft/agents/hosting/core/storage/storage_test_utils.py

@@ -92,22 +92,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):
@@ -133,10 +117,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 +143,7 @@ async def equals(self, other) -> bool:
 
 
 class StorageTestsCommon(ABC):
+    """Common fixtures for Storage implementations."""
 
     KEY_LIST = [
         "f",
@@ -211,8 +200,16 @@ 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) -> 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.
+        """
         raise NotImplementedError("Subclasses must implement this")
 
     @pytest.mark.asyncio
@@ -446,9 +443,64 @@ 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()
             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
+
+
+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)

libraries/microsoft-agents-hosting-core/tests/test_memory_storage.py

@@ -1,25 +1,11 @@
 from microsoft.agents.hosting.core.storage.memory_storage import MemoryStorage
-from microsoft.agents.hosting.core.storage.storage_test_utils import (
-    CRUDStorageTests,
-    StorageMock,
-)
+from microsoft.agents.hosting.core.storage.storage_test_utils import CRUDStorageTests
 
 
-class MemoryStorageMock(StorageMock):
-
-    def __init__(self, initial_data: dict = None):
-
+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)

libraries/microsoft-agents-storage-blob/tests/test_blob_storage.py

@@ -12,7 +12,6 @@
 
 from microsoft.agents.hosting.core.storage.storage_test_utils import (
     CRUDStorageTests,
-    StorageMock,
     StorageBaseline,
     MockStoreItem,
     MockStoreItemB,
@@ -69,15 +68,6 @@ async def blob_storage():
     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 +80,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 +94,26 @@ 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)) == {}
+        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
+        )
+        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))[
+            "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):
 
@@ -183,5 +193,4 @@ async def test_blob_storage_flow_existing_container_and_persistence(self):
             == initial_data["1230"]
         )
 
-        # teardown
         await container_client.delete_container()

libraries/microsoft-agents-storage-cosmos/microsoft/agents/storage/cosmos/__init__.py

@@ -0,0 +1,7 @@
+from .cosmos_db_storage import CosmosDBStorage
+from .cosmos_db_storage_config import CosmosDBStorageConfig
+
+__all__ = [
+    "CosmosDBStorage",
+    "CosmosDBStorageConfig",
+]