microsoft-agents-storage-cosmos 0.3.2__tar.gz → 0.4.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (17) hide show
  1. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/PKG-INFO +2 -2
  2. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/microsoft_agents_storage_cosmos.egg-info/PKG-INFO +2 -2
  3. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/microsoft_agents_storage_cosmos.egg-info/SOURCES.txt +1 -4
  4. microsoft_agents_storage_cosmos-0.4.0/microsoft_agents_storage_cosmos.egg-info/requires.txt +3 -0
  5. microsoft_agents_storage_cosmos-0.3.2/microsoft_agents_storage_cosmos.egg-info/requires.txt +0 -3
  6. microsoft_agents_storage_cosmos-0.3.2/tests/test_cosmos_db_config.py +0 -245
  7. microsoft_agents_storage_cosmos-0.3.2/tests/test_cosmos_db_storage.py +0 -299
  8. microsoft_agents_storage_cosmos-0.3.2/tests/test_key_ops.py +0 -250
  9. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/microsoft_agents/storage/cosmos/__init__.py +0 -0
  10. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/microsoft_agents/storage/cosmos/cosmos_db_storage.py +0 -0
  11. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/microsoft_agents/storage/cosmos/cosmos_db_storage_config.py +0 -0
  12. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/microsoft_agents/storage/cosmos/key_ops.py +0 -0
  13. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/microsoft_agents_storage_cosmos.egg-info/dependency_links.txt +0 -0
  14. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/microsoft_agents_storage_cosmos.egg-info/top_level.txt +0 -0
  15. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/pyproject.toml +0 -0
  16. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/setup.cfg +0 -0
  17. {microsoft_agents_storage_cosmos-0.3.2 → microsoft_agents_storage_cosmos-0.4.0}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: microsoft-agents-storage-cosmos
3
- Version: 0.3.2
3
+ Version: 0.4.0
4
4
  Summary: A Cosmos DB storage library for Microsoft Agents
5
5
  Author: Microsoft Corporation
6
6
  Project-URL: Homepage, https://github.com/microsoft/Agents
@@ -8,7 +8,7 @@ Classifier: Programming Language :: Python :: 3
8
8
  Classifier: License :: OSI Approved :: MIT License
9
9
  Classifier: Operating System :: OS Independent
10
10
  Requires-Python: >=3.9
11
- Requires-Dist: microsoft-agents-hosting-core==0.3.2
11
+ Requires-Dist: microsoft-agents-hosting-core==0.4.0
12
12
  Requires-Dist: azure-core
13
13
  Requires-Dist: azure-cosmos
14
14
  Dynamic: requires-dist
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: microsoft-agents-storage-cosmos
3
- Version: 0.3.2
3
+ Version: 0.4.0
4
4
  Summary: A Cosmos DB storage library for Microsoft Agents
5
5
  Author: Microsoft Corporation
6
6
  Project-URL: Homepage, https://github.com/microsoft/Agents
@@ -8,7 +8,7 @@ Classifier: Programming Language :: Python :: 3
8
8
  Classifier: License :: OSI Approved :: MIT License
9
9
  Classifier: Operating System :: OS Independent
10
10
  Requires-Python: >=3.9
11
- Requires-Dist: microsoft-agents-hosting-core==0.3.2
11
+ Requires-Dist: microsoft-agents-hosting-core==0.4.0
12
12
  Requires-Dist: azure-core
13
13
  Requires-Dist: azure-cosmos
14
14
  Dynamic: requires-dist
@@ -8,7 +8,4 @@ microsoft_agents_storage_cosmos.egg-info/PKG-INFO
8
8
  microsoft_agents_storage_cosmos.egg-info/SOURCES.txt
9
9
  microsoft_agents_storage_cosmos.egg-info/dependency_links.txt
10
10
  microsoft_agents_storage_cosmos.egg-info/requires.txt
11
- microsoft_agents_storage_cosmos.egg-info/top_level.txt
12
- tests/test_cosmos_db_config.py
13
- tests/test_cosmos_db_storage.py
14
- tests/test_key_ops.py
11
+ microsoft_agents_storage_cosmos.egg-info/top_level.txt
@@ -0,0 +1,3 @@
1
+ microsoft-agents-hosting-core==0.4.0
2
+ azure-core
3
+ azure-cosmos
@@ -1,3 +0,0 @@
1
- microsoft-agents-hosting-core==0.3.2
2
- azure-core
3
- azure-cosmos
@@ -1,245 +0,0 @@
1
- import json
2
- import pytest
3
-
4
- from microsoft_agents.storage.cosmos import CosmosDBStorageConfig
5
-
6
- # thank you AI, again
7
-
8
-
9
- @pytest.fixture()
10
- def valid_config():
11
- """Fixture providing a valid CosmosDBStorageConfig for tests"""
12
- return CosmosDBStorageConfig(
13
- cosmos_db_endpoint="https://localhost:8081",
14
- auth_key=(
15
- "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGG"
16
- "yPMbIZnqyMsEcaGQy67XIw/Jw=="
17
- ),
18
- database_id="test-db",
19
- container_id="bot-storage",
20
- )
21
-
22
-
23
- @pytest.fixture()
24
- def minimal_config():
25
- """Fixture providing a minimal CosmosDBStorageConfig for tests"""
26
- return CosmosDBStorageConfig()
27
-
28
-
29
- @pytest.fixture()
30
- def config_with_options():
31
- """Fixture providing a CosmosDBStorageConfig with all options for tests"""
32
- return CosmosDBStorageConfig(
33
- cosmos_db_endpoint="https://test.documents.azure.com:443/",
34
- auth_key="test_key",
35
- database_id="test_db",
36
- container_id="test_container",
37
- cosmos_client_options={"connection_policy": "test"},
38
- container_throughput=800,
39
- key_suffix="_test",
40
- compatibility_mode=False,
41
- )
42
-
43
-
44
- class TestCosmosDBStorageConfig:
45
-
46
- def test_constructor_with_parameters(self):
47
- """Test creating config with direct parameters"""
48
- config = CosmosDBStorageConfig(
49
- cosmos_db_endpoint="https://test.documents.azure.com:443/",
50
- auth_key="test_key",
51
- database_id="test_db",
52
- container_id="test_container",
53
- container_throughput=800,
54
- key_suffix="_test",
55
- compatibility_mode=False,
56
- )
57
-
58
- assert config.cosmos_db_endpoint == "https://test.documents.azure.com:443/"
59
- assert config.auth_key == "test_key"
60
- assert config.database_id == "test_db"
61
- assert config.container_id == "test_container"
62
- assert config.container_throughput == 800
63
- assert config.key_suffix == "_test"
64
- assert config.compatibility_mode is False
65
- assert config.cosmos_client_options == {}
66
- assert config.credential is None
67
-
68
- def test_constructor_with_defaults(self):
69
- """Test creating config with default values"""
70
- config = CosmosDBStorageConfig()
71
-
72
- assert config.cosmos_db_endpoint == ""
73
- assert config.auth_key == ""
74
- assert config.database_id == ""
75
- assert config.container_id == ""
76
- assert config.container_throughput == 400 # Default value
77
- assert config.key_suffix == ""
78
- assert config.compatibility_mode is False
79
- assert config.cosmos_client_options == {}
80
- assert config.credential is None
81
-
82
- def test_from_file(self, tmp_path):
83
- """Test creating config from JSON file"""
84
- config_file_path = tmp_path / "cosmos_config.json"
85
-
86
- config_data = {
87
- "cosmos_db_endpoint": "https://localhost:8081",
88
- "auth_key": "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==",
89
- "database_id": "test-db",
90
- "container_id": "bot-storage",
91
- "container_throughput": 600,
92
- "key_suffix": "_file",
93
- "compatibility_mode": True,
94
- "cosmos_client_options": {"connection_policy": "test"},
95
- }
96
-
97
- with open(config_file_path, "w") as f:
98
- json.dump(config_data, f)
99
-
100
- config = CosmosDBStorageConfig(filename=str(config_file_path))
101
-
102
- assert config.cosmos_db_endpoint == "https://localhost:8081"
103
- assert (
104
- config.auth_key
105
- == "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="
106
- )
107
- assert config.database_id == "test-db"
108
- assert config.container_id == "bot-storage"
109
- assert config.container_throughput == 600
110
- assert config.key_suffix == "_file"
111
- assert config.compatibility_mode is True
112
- assert config.cosmos_client_options == {"connection_policy": "test"}
113
-
114
- def test_parameter_override_file(self, tmp_path):
115
- """Test that constructor parameters override file values"""
116
- config_file_path = tmp_path / "cosmos_config.json"
117
-
118
- with open(config_file_path, "w") as f:
119
- json.dump(
120
- {
121
- "cosmos_db_endpoint": "https://file-endpoint.com",
122
- "auth_key": "file_key",
123
- "database_id": "file_db",
124
- },
125
- f,
126
- )
127
-
128
- config = CosmosDBStorageConfig(
129
- cosmos_db_endpoint="https://param-endpoint.com",
130
- auth_key="param_key",
131
- filename=str(config_file_path),
132
- )
133
-
134
- # Parameters should override file values
135
- assert config.cosmos_db_endpoint == "https://param-endpoint.com"
136
- assert config.auth_key == "param_key"
137
- # File value should be used when parameter not provided
138
- assert config.database_id == "file_db"
139
-
140
- def test_validation_success(self):
141
- """Test successful validation with all required fields"""
142
- config = CosmosDBStorageConfig(
143
- cosmos_db_endpoint="https://test.documents.azure.com:443/",
144
- auth_key="test_key",
145
- database_id="test_db",
146
- container_id="test_container",
147
- )
148
-
149
- # Should not raise any exception
150
- CosmosDBStorageConfig.validate_cosmos_db_config(config)
151
-
152
- def test_validation_missing_config(self):
153
- """Test validation with None config"""
154
- with pytest.raises(ValueError):
155
- CosmosDBStorageConfig.validate_cosmos_db_config(None)
156
-
157
- def test_validation_missing_endpoint(self):
158
- """Test validation with missing cosmos_db_endpoint"""
159
- config = CosmosDBStorageConfig(
160
- auth_key="test_key", database_id="test_db", container_id="test_container"
161
- )
162
- with pytest.raises(ValueError):
163
- CosmosDBStorageConfig.validate_cosmos_db_config(config)
164
-
165
- def test_validation_missing_auth_key(self):
166
- """Test validation with missing auth_key"""
167
- config = CosmosDBStorageConfig(
168
- cosmos_db_endpoint="https://test.documents.azure.com:443/",
169
- database_id="test_db",
170
- container_id="test_container",
171
- )
172
- with pytest.raises(ValueError):
173
- CosmosDBStorageConfig.validate_cosmos_db_config(config)
174
-
175
- def test_validation_missing_database_id(self):
176
- """Test validation with missing database_id"""
177
- config = CosmosDBStorageConfig(
178
- cosmos_db_endpoint="https://test.documents.azure.com:443/",
179
- auth_key="test_key",
180
- container_id="test_container",
181
- )
182
- with pytest.raises(ValueError):
183
- CosmosDBStorageConfig.validate_cosmos_db_config(config)
184
-
185
- def test_validation_missing_container_id(self):
186
- """Test validation with missing container_id"""
187
- config = CosmosDBStorageConfig(
188
- cosmos_db_endpoint="https://test.documents.azure.com:443/",
189
- auth_key="test_key",
190
- database_id="test_db",
191
- )
192
- with pytest.raises(ValueError):
193
- CosmosDBStorageConfig.validate_cosmos_db_config(config)
194
-
195
- def test_validation_suffix_with_compatibility_mode(self):
196
- """Test validation fails when using suffix with compatibility mode"""
197
- config = CosmosDBStorageConfig(
198
- cosmos_db_endpoint="https://test.documents.azure.com:443/",
199
- auth_key="test_key",
200
- database_id="test_db",
201
- container_id="test_container",
202
- key_suffix="_test",
203
- compatibility_mode=True,
204
- )
205
- with pytest.raises(ValueError):
206
- CosmosDBStorageConfig.validate_cosmos_db_config(config)
207
-
208
- def test_validation_invalid_suffix_characters(self):
209
- """Test validation fails with invalid characters in suffix"""
210
- config = CosmosDBStorageConfig(
211
- cosmos_db_endpoint="https://test.documents.azure.com:443/",
212
- auth_key="test_key",
213
- database_id="test_db",
214
- container_id="test_container",
215
- key_suffix="invalid/suffix\\with?bad#chars",
216
- compatibility_mode=False,
217
- )
218
- with pytest.raises(ValueError, match="Cannot use invalid Row Key characters"):
219
- CosmosDBStorageConfig.validate_cosmos_db_config(config)
220
-
221
- def test_validation_valid_suffix(self):
222
- """Test validation succeeds with valid suffix"""
223
- config = CosmosDBStorageConfig(
224
- cosmos_db_endpoint="https://test.documents.azure.com:443/",
225
- auth_key="test_key",
226
- database_id="test_db",
227
- container_id="test_container",
228
- key_suffix="valid_suffix_123",
229
- compatibility_mode=False,
230
- )
231
- # Should not raise any exception
232
- CosmosDBStorageConfig.validate_cosmos_db_config(config)
233
-
234
- def test_cosmos_client_options(self):
235
- """Test cosmos_client_options handling"""
236
- options = {"connection_policy": "test", "consistency_level": "strong"}
237
- config = CosmosDBStorageConfig(cosmos_client_options=options)
238
- assert config.cosmos_client_options == options
239
-
240
- def test_credential_parameter(self):
241
- """Test credential parameter handling"""
242
- # Mock credential (in real usage this would be a TokenCredential instance)
243
- mock_credential = object() # Placeholder for actual TokenCredential
244
- config = CosmosDBStorageConfig(credential=mock_credential)
245
- assert config.credential is mock_credential
@@ -1,299 +0,0 @@
1
- # Copyright (c) Microsoft Corporation. All rights reserved.
2
- # Licensed under the MIT License.
3
-
4
- import gc
5
-
6
- import pytest
7
- import pytest_asyncio
8
-
9
- from azure.cosmos import documents
10
- from azure.cosmos.aio import CosmosClient
11
- from azure.cosmos.exceptions import CosmosResourceNotFoundError
12
-
13
- from microsoft_agents.storage.cosmos import CosmosDBStorage, CosmosDBStorageConfig
14
- from microsoft_agents.storage.cosmos.key_ops import sanitize_key
15
-
16
- from microsoft_agents.hosting.core.storage.storage_test_utils import (
17
- QuickCRUDStorageTests,
18
- MockStoreItem,
19
- MockStoreItemB,
20
- StorageBaseline,
21
- )
22
-
23
- EMULATOR_RUNNING = False
24
-
25
-
26
- def create_config(compat_mode):
27
- return CosmosDBStorageConfig(
28
- cosmos_db_endpoint="https://localhost:8081",
29
- auth_key=(
30
- "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGG"
31
- "yPMbIZnqyMsEcaGQy67XIw/Jw=="
32
- ),
33
- database_id="test-db",
34
- container_id="bot-storage",
35
- compatibility_mode=compat_mode,
36
- container_throughput=800,
37
- )
38
-
39
-
40
- @pytest.fixture
41
- def config():
42
- return create_config(compat_mode=False)
43
-
44
-
45
- async def create_cosmos_env(config, compat_mode=False, existing=False):
46
- """Creates the Cosmos DB environment for testing.
47
-
48
- If existing is False, creates a new database and container, deleting any
49
- existing ones with the same name. If existing is True, creates the database
50
- and container if they do not already exist."""
51
-
52
- cosmos_client = CosmosClient(
53
- config.cosmos_db_endpoint,
54
- config.auth_key,
55
- )
56
-
57
- if not existing:
58
- try:
59
- await cosmos_client.delete_database(config.database_id)
60
- except Exception:
61
- pass
62
- database = await cosmos_client.create_database(id=config.database_id)
63
-
64
- try:
65
- await database.delete_container(config.container_id)
66
- except Exception:
67
- pass
68
-
69
- partition_key = {
70
- "paths": ["/_partitionKey"] if compat_mode else ["/id"],
71
- "kind": documents.PartitionKind.Hash,
72
- }
73
- container_client = await database.create_container(
74
- id=config.container_id,
75
- partition_key=partition_key,
76
- offer_throughput=config.container_throughput,
77
- )
78
- else:
79
- database = await cosmos_client.create_database_if_not_exists(
80
- id=config.database_id
81
- )
82
- container_client = database.get_container_client(config.container_id)
83
-
84
- return container_client
85
-
86
-
87
- async def cosmos_db_storage_instance(compat_mode=False, existing=False):
88
- config = create_config(compat_mode)
89
- container_client = await create_cosmos_env(
90
- config, compat_mode=compat_mode, existing=existing
91
- )
92
- storage = CosmosDBStorage(config)
93
- return storage, container_client
94
-
95
-
96
- @pytest_asyncio.fixture()
97
- async def cosmos_db_storage():
98
- storage, _ = await cosmos_db_storage_instance()
99
- return storage
100
-
101
-
102
- @pytest.mark.asyncio
103
- @pytest.mark.parametrize("test_require_compat", [True, False])
104
- @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.")
105
- async def test_cosmos_db_storage_flow_existing_container_and_persistence(
106
- test_require_compat,
107
- ):
108
-
109
- config = create_config(compat_mode=test_require_compat)
110
- container_client = await create_cosmos_env(config)
111
-
112
- initial_data = {
113
- "__some_key": MockStoreItem({"id": "item2", "value": "data2"}),
114
- "?test": MockStoreItem({"id": "?test", "value": "data1"}),
115
- "!another_key": MockStoreItem({"id": "item3", "value": "data3"}),
116
- "1230": MockStoreItemB({"id": "item8", "value": "data"}, False),
117
- "key-with-dash": MockStoreItem({"id": "item4", "value": "data"}),
118
- "key.with.dot": MockStoreItem({"id": "item5", "value": "data"}),
119
- "key/with/slash": MockStoreItem({"id": "item6", "value": "data"}),
120
- "another key": MockStoreItemB({"id": "item7", "value": "data"}, True),
121
- }
122
-
123
- baseline_storage = StorageBaseline(initial_data)
124
-
125
- for key, value in initial_data.items():
126
- doc = {
127
- "id": sanitize_key(
128
- key,
129
- config.key_suffix,
130
- test_require_compat,
131
- ),
132
- "realId": key,
133
- "document": value.store_item_to_json(),
134
- }
135
- await container_client.upsert_item(body=doc)
136
-
137
- storage = CosmosDBStorage(config)
138
- assert await baseline_storage.equals(storage)
139
- assert (
140
- await storage.read(["1230", "another key"], target_cls=MockStoreItemB)
141
- ) == baseline_storage.read(["1230", "another key"])
142
-
143
- changes = {
144
- "?test": MockStoreItem({"id": "?test", "value": "data1_changed"}),
145
- "__some_key": MockStoreItem({"id": "item2", "value": "data2_changed"}),
146
- "new_item": MockStoreItem({"id": "new_item", "value": "new_data"}),
147
- }
148
-
149
- baseline_storage.write(changes)
150
- await storage.write(changes)
151
-
152
- baseline_storage.delete(["!another_key", "?test"])
153
- await storage.delete(["!another_key", "?test"])
154
- assert await baseline_storage.equals(storage)
155
-
156
- del storage
157
- gc.collect()
158
- storage = CosmosDBStorage(config)
159
-
160
- escaped_key = storage._sanitize("?test")
161
- with pytest.raises(CosmosResourceNotFoundError):
162
- await container_client.read_item(
163
- escaped_key, storage._get_partition_key(escaped_key)
164
- )
165
-
166
- escaped_key = storage._sanitize("1230")
167
- item = (
168
- await container_client.read_item(
169
- escaped_key, storage._get_partition_key(escaped_key)
170
- )
171
- ).get("document")
172
- assert MockStoreItemB.from_json_to_store_item(item) == initial_data["1230"]
173
-
174
-
175
- @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.")
176
- class TestCosmosDBStorage(QuickCRUDStorageTests):
177
-
178
- def get_compat_mode(self):
179
- return False
180
-
181
- async def storage(self, initial_data=None, existing=False):
182
- storage, _ = await cosmos_db_storage_instance(
183
- compat_mode=self.get_compat_mode(), existing=existing
184
- )
185
- if initial_data:
186
- await storage.write(initial_data)
187
- return storage
188
-
189
- @pytest.mark.asyncio
190
- async def test_initialize(self, cosmos_db_storage):
191
- await cosmos_db_storage.initialize()
192
- await cosmos_db_storage.initialize()
193
- await cosmos_db_storage.write(
194
- {"some_Key": MockStoreItem({"id": "123", "data": "value"})}
195
- )
196
- await cosmos_db_storage.initialize()
197
- assert (
198
- await cosmos_db_storage.read(["some_Key"], target_cls=MockStoreItem)
199
- ) == {"some_Key": MockStoreItem({"id": "123", "data": "value"})}
200
-
201
- @pytest.mark.asyncio
202
- async def test_external_change_is_visible(self):
203
- cosmos_storage, container_client = await cosmos_db_storage_instance()
204
- assert (await cosmos_storage.read(["key"], target_cls=MockStoreItem)) == {}
205
- assert (await cosmos_storage.read(["key2"], target_cls=MockStoreItem)) == {}
206
- await container_client.upsert_item(
207
- {
208
- "id": "key",
209
- "realId": "key",
210
- "document": {"id": "key", "value": "data"},
211
- "partitionKey": "",
212
- }
213
- )
214
- await container_client.upsert_item(
215
- {
216
- "id": "key2",
217
- "realId": "key2",
218
- "document": {"id": "key2", "value": "new_val"},
219
- "partitionKey": "",
220
- }
221
- )
222
- assert (await cosmos_storage.read(["key"], target_cls=MockStoreItem))[
223
- "key"
224
- ] == MockStoreItem({"id": "key", "value": "data"})
225
- assert (await cosmos_storage.read(["key2"], target_cls=MockStoreItem))[
226
- "key2"
227
- ] == MockStoreItem({"id": "key2", "value": "new_val"})
228
-
229
-
230
- @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.")
231
- class TestCosmosDBStorageWithCompat(TestCosmosDBStorage):
232
- def get_compat_mode(self):
233
- return True
234
-
235
-
236
- @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.")
237
- class TestCosmosDBStorageInit:
238
-
239
- def test_raises_error_when_no_endpoint_provided(self, config):
240
- config.cosmos_db_endpoint = None
241
- with pytest.raises(ValueError):
242
- CosmosDBStorage(config)
243
-
244
- def test_raises_error_when_no_auth_key_provided(self, config):
245
- config.auth_key = None
246
- with pytest.raises(ValueError):
247
- CosmosDBStorage(config)
248
-
249
- def test_raises_error_when_suffix_provided_but_compat(self, config):
250
- config.auth_key = None
251
- config.compatibility_mode = True
252
- with pytest.raises(ValueError):
253
- CosmosDBStorage(config)
254
-
255
- def test_raises_error_when_no_database_id_provided(self, config):
256
- config.database_id = None
257
- with pytest.raises(ValueError):
258
- CosmosDBStorage(config)
259
-
260
- def test_raises_error_when_no_container_id_provided(self, config):
261
- config.container_id = None
262
- with pytest.raises(ValueError):
263
- CosmosDBStorage(config)
264
-
265
- @pytest.mark.asyncio
266
- @pytest.mark.parametrize("compat_mode", [True, False])
267
- async def test_raises_error_different_partition_key(self, compat_mode):
268
- config = create_config(compat_mode=compat_mode)
269
- await create_cosmos_env(config, compat_mode=compat_mode)
270
- storage = CosmosDBStorage(config)
271
-
272
- with pytest.raises(Exception):
273
-
274
- cosmos_client = CosmosClient(
275
- config.cosmos_db_endpoint,
276
- config.auth_key,
277
- )
278
- try:
279
- await cosmos_client.delete_database(config.database_id)
280
- except Exception:
281
- pass
282
- database = await cosmos_client.create_database(id=config.database_id)
283
-
284
- try:
285
- await database.delete_container(config.container_id)
286
- except Exception:
287
- pass
288
-
289
- partition_key = {
290
- "paths": ["/fake_part_key"],
291
- "kind": documents.PartitionKind.Hash,
292
- }
293
- container_client = await database.create_container(
294
- id=config.container_id,
295
- partition_key=partition_key,
296
- offer_throughput=config.container_throughput,
297
- )
298
- storage = CosmosDBStorage(config)
299
- await storage.initialize()
@@ -1,250 +0,0 @@
1
- import hashlib
2
- import pytest
3
- from microsoft_agents.storage.cosmos.key_ops import truncate_key, sanitize_key
4
-
5
- # thank you AI
6
-
7
-
8
- @pytest.mark.parametrize(
9
- "input_key,expected",
10
- [
11
- ("validKey123", "validKey123"),
12
- ("simple", "simple"),
13
- ("CamelCase", "CamelCase"),
14
- ("under_score", "under_score"),
15
- ("with-dash", "with-dash"),
16
- ("with.dot", "with.dot"),
17
- ],
18
- )
19
- def test_sanitize_key_simple(input_key, expected):
20
- assert sanitize_key(input_key) == expected
21
-
22
-
23
- @pytest.mark.parametrize(
24
- "input_key,expected",
25
- [
26
- ("key\\value", "key*92value"),
27
- ("key?value", "key*63value"),
28
- ("key/value", "key*47value"),
29
- ("key#value", "key*35value"),
30
- ("key\tvalue", "key*9value"),
31
- ("key\nvalue", "key*10value"),
32
- ("key\rvalue", "key*13value"),
33
- ("key*value", "key*42value"),
34
- ],
35
- )
36
- def test_sanitize_key_forbidden_chars(input_key, expected):
37
- assert sanitize_key(input_key) == expected
38
-
39
-
40
- @pytest.mark.parametrize(
41
- "input_key,expected",
42
- [
43
- ("key/with\\many?bad#chars", "key*47with*92many*63bad*35chars"),
44
- ("a\\b/c?d#e\tf\ng\rh*i", "a*92b*47c*63d*35e*9f*10g*13h*42i"),
45
- ("key/with\\many?bad#chars", "key*47with*92many*63bad*35chars"),
46
- ],
47
- )
48
- def test_sanitize_key_multiple_forbidden_chars(input_key, expected):
49
- assert sanitize_key(input_key) == expected
50
-
51
-
52
- def test_sanitize_key_with_long_key_with_forbidden_chars():
53
- long_key = "a?2/!@\t3." * 100 # Create a long key
54
- sanitized = sanitize_key(long_key)
55
- assert len(sanitized) <= 255 # Should be truncated
56
- # Ensure forbidden characters are replaced
57
- assert "?" not in sanitized
58
- assert "/" not in sanitized
59
- assert "\t" not in sanitized
60
-
61
-
62
- def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix():
63
- long_key = "a?2/!@\t3." * 100 # Create a long key
64
- sanitized = sanitize_key(long_key, key_suffix="_suff?#*")
65
- assert len(sanitized) <= 255 # Should be truncated
66
- # Ensure forbidden characters are replaced
67
- assert "?" not in sanitized
68
- assert "/" not in sanitized
69
- assert "#" not in sanitized
70
-
71
-
72
- def test_sanitize_key_with_long_key_with_forbidden_chars_with_suffix_compat_mode():
73
- long_key = "a?2/!@\t3." * 100 # Create a long key
74
- sanitized = sanitize_key(long_key, key_suffix="_suff?#*", compatibility_mode=True)
75
- assert len(sanitized) <= 255 # Should be truncated
76
- # Ensure forbidden characters are replaced
77
- assert "?" not in sanitized
78
- assert "/" not in sanitized
79
- assert "#" not in sanitized
80
-
81
-
82
- @pytest.mark.parametrize(
83
- "input_key,expected",
84
- [
85
- ("", ""),
86
- (" ", " "),
87
- ],
88
- )
89
- def test_sanitize_key_empty_and_whitespace(input_key, expected):
90
- assert sanitize_key(input_key) == expected
91
-
92
-
93
- @pytest.mark.parametrize(
94
- "input_key,suffix,expected",
95
- [
96
- ("key", "_suffix", "key_suffix"),
97
- ("test", "123", "test123"),
98
- ("key/value", "_clean", "key*47value_clean"),
99
- ("", "_suffix", "_suffix"),
100
- ],
101
- )
102
- def test_sanitize_key_with_suffix(input_key, suffix, expected):
103
- assert sanitize_key(input_key, key_suffix=suffix) == expected
104
-
105
-
106
- def test_sanitize_key_suffix_with_truncation():
107
- long_key = "a" * 250
108
- suffix = "_suffix"
109
- result = sanitize_key(long_key, key_suffix=suffix, compatibility_mode=True)
110
- assert len(result) <= 255
111
- assert (
112
- result.endswith(suffix) or len(result) == 255
113
- ) # Either has suffix or was truncated
114
-
115
-
116
- def test_sanitize_key_truncation_compatibility_mode():
117
- long_key = "a" * 300
118
- result = sanitize_key(long_key, compatibility_mode=True)
119
- assert len(result) <= 255
120
-
121
- # Should contain hash when truncated
122
- very_long_key = "b" * 500
123
- result2 = sanitize_key(very_long_key, compatibility_mode=True)
124
- assert len(result2) == 255
125
-
126
-
127
- def test_sanitize_key_no_truncation():
128
- long_key = "a" * 300
129
- result = sanitize_key(long_key, compatibility_mode=False)
130
- assert result == long_key # Should be unchanged
131
- assert len(result) == 300
132
-
133
-
134
- @pytest.mark.parametrize(
135
- "input_key,expected",
136
- [
137
- ("short", "short"),
138
- ("a" * 254, "a" * 254),
139
- ("a" * 255, "a" * 255),
140
- ],
141
- )
142
- def test_truncate_key_short_strings(input_key, expected):
143
- assert truncate_key(input_key) == expected
144
-
145
-
146
- def test_truncate_key_long_strings():
147
- long_key = "a" * 300
148
- result = truncate_key(long_key)
149
- assert len(result) == 255
150
-
151
- # Result should end with SHA256 hash
152
- expected_hash = hashlib.sha256(long_key.encode("utf-8")).hexdigest()
153
- assert result.endswith(expected_hash)
154
-
155
- # First part should be original key truncated
156
- expected_prefix_len = 255 - len(expected_hash)
157
- assert result.startswith("a" * expected_prefix_len)
158
-
159
-
160
- @pytest.mark.parametrize(
161
- "input_key,compatibility_mode,expected_unchanged",
162
- [
163
- ("a" * 300, False, True), # Should be unchanged
164
- ("x" * 1000, False, True), # Should be unchanged
165
- (
166
- "key/with\\special?chars#and\ttabs\nand\rmore*",
167
- False,
168
- True,
169
- ), # Should be unchanged
170
- ],
171
- )
172
- def test_truncate_key_compatibility_mode_disabled(
173
- input_key, compatibility_mode, expected_unchanged
174
- ):
175
- result = truncate_key(input_key, compatibility_mode=compatibility_mode)
176
- if expected_unchanged:
177
- assert result == input_key
178
-
179
-
180
- @pytest.mark.parametrize(
181
- "input_key,expected_length",
182
- [
183
- ("a" * 255, 255),
184
- ("a" * 256, 255),
185
- ],
186
- )
187
- def test_truncate_key_exact_and_over_limit(input_key, expected_length):
188
- result = truncate_key(input_key)
189
- assert len(result) == expected_length
190
-
191
- if len(input_key) == 255:
192
- assert result == input_key
193
- else:
194
- assert result != input_key
195
-
196
-
197
- def test_truncate_key_hash_consistency():
198
- long_key = "consistent_test_key_" * 20 # > 255 chars
199
- result1 = truncate_key(long_key)
200
- result2 = truncate_key(long_key)
201
- assert result1 == result2
202
- assert len(result1) == 255
203
-
204
-
205
- @pytest.mark.parametrize(
206
- "key1,key2",
207
- [
208
- ("a" * 300, "b" * 300),
209
- ("consistent_test_key_" * 20, "different_test_key_" * 20),
210
- ],
211
- )
212
- def test_truncate_key_different_inputs_different_outputs(key1, key2):
213
- result1 = truncate_key(key1)
214
- result2 = truncate_key(key2)
215
- assert result1 != result2
216
- assert len(result1) == len(result2) == 255
217
-
218
-
219
- def test_sanitize_key_integration():
220
- # Key with forbidden chars that will be long after sanitization + suffix
221
- base_key = "test/key\\with?many#forbidden\tchars\nand\rmore*" * 10
222
- suffix = "_integration_test"
223
-
224
- result = sanitize_key(base_key, key_suffix=suffix, compatibility_mode=True)
225
-
226
- # Should be sanitized and truncated
227
- assert len(result) <= 255
228
- assert "*47" in result or "*92" in result # Contains sanitized chars
229
-
230
- # Test without truncation
231
- result_no_trunc = sanitize_key(
232
- base_key, key_suffix=suffix, compatibility_mode=False
233
- )
234
- assert (
235
- "*47" in result_no_trunc or "*92" in result_no_trunc
236
- ) # Contains sanitized chars
237
- assert result_no_trunc.endswith(suffix)
238
-
239
-
240
- @pytest.mark.parametrize(
241
- "input_key,expected",
242
- [
243
- ("key_ñ_测试", "key_ñ_测试"),
244
- ("123456789", "123456789"),
245
- ("MyKey/WithSlash", "MyKey*47WithSlash"),
246
- ],
247
- )
248
- def test_edge_cases(input_key, expected):
249
- result = sanitize_key(input_key)
250
- assert result == expected