Spaces:
Running
Running
| import pytest | |
| from chromadb.api.client import AdminClient, Client | |
| from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT | |
| def test_database_tenant_collections(client: Client) -> None: | |
| client.reset() | |
| # Create a new database in the default tenant | |
| admin_client = AdminClient.from_system(client._system) | |
| admin_client.create_database("test_db") | |
| # Create collections in this new database | |
| client.set_tenant(tenant=DEFAULT_TENANT, database="test_db") | |
| client.create_collection("collection", metadata={"database": "test_db"}) | |
| # Create collections in the default database | |
| client.set_tenant(tenant=DEFAULT_TENANT, database=DEFAULT_DATABASE) | |
| client.create_collection("collection", metadata={"database": DEFAULT_DATABASE}) | |
| # List collections in the default database | |
| collections = client.list_collections() | |
| assert len(collections) == 1 | |
| assert collections[0].name == "collection" | |
| assert collections[0].metadata == {"database": DEFAULT_DATABASE} | |
| # List collections in the new database | |
| client.set_tenant(tenant=DEFAULT_TENANT, database="test_db") | |
| collections = client.list_collections() | |
| assert len(collections) == 1 | |
| assert collections[0].metadata == {"database": "test_db"} | |
| # Update the metadata in both databases to different values | |
| client.set_tenant(tenant=DEFAULT_TENANT, database=DEFAULT_DATABASE) | |
| client.list_collections()[0].modify(metadata={"database": "default2"}) | |
| client.set_tenant(tenant=DEFAULT_TENANT, database="test_db") | |
| client.list_collections()[0].modify(metadata={"database": "test_db2"}) | |
| # Validate that the metadata was updated | |
| client.set_tenant(tenant=DEFAULT_TENANT, database=DEFAULT_DATABASE) | |
| collections = client.list_collections() | |
| assert len(collections) == 1 | |
| assert collections[0].metadata == {"database": "default2"} | |
| client.set_tenant(tenant=DEFAULT_TENANT, database="test_db") | |
| collections = client.list_collections() | |
| assert len(collections) == 1 | |
| assert collections[0].metadata == {"database": "test_db2"} | |
| # Delete the collections and make sure databases are isolated | |
| client.set_tenant(tenant=DEFAULT_TENANT, database=DEFAULT_DATABASE) | |
| client.delete_collection("collection") | |
| collections = client.list_collections() | |
| assert len(collections) == 0 | |
| client.set_tenant(tenant=DEFAULT_TENANT, database="test_db") | |
| collections = client.list_collections() | |
| assert len(collections) == 1 | |
| client.delete_collection("collection") | |
| collections = client.list_collections() | |
| assert len(collections) == 0 | |
| def test_database_collections_add(client: Client) -> None: | |
| client.reset() | |
| # Create a new database in the default tenant | |
| admin_client = AdminClient.from_system(client._system) | |
| admin_client.create_database("test_db") | |
| # Create collections in this new database | |
| client.set_database(database="test_db") | |
| coll_new = client.create_collection("collection_new") | |
| # Create collections in the default database | |
| client.set_database(database=DEFAULT_DATABASE) | |
| coll_default = client.create_collection("collection_default") | |
| records_new = { | |
| "ids": ["a", "b", "c"], | |
| "embeddings": [[1.0, 2.0, 3.0] for _ in range(3)], | |
| "documents": ["a", "b", "c"], | |
| } | |
| records_default = { | |
| "ids": ["c", "d", "e"], | |
| "embeddings": [[4.0, 5.0, 6.0] for _ in range(3)], | |
| "documents": ["c", "d", "e"], | |
| } | |
| # Add to the new coll | |
| coll_new.add(**records_new) # type: ignore | |
| # Add to the default coll | |
| coll_default.add(**records_default) # type: ignore | |
| # Make sure the collections are isolated | |
| res = coll_new.get(include=["embeddings", "documents"]) | |
| assert res["ids"] == records_new["ids"] | |
| assert res["embeddings"] == records_new["embeddings"] | |
| assert res["documents"] == records_new["documents"] | |
| res = coll_default.get(include=["embeddings", "documents"]) | |
| assert res["ids"] == records_default["ids"] | |
| assert res["embeddings"] == records_default["embeddings"] | |
| assert res["documents"] == records_default["documents"] | |
| def test_tenant_collections_add(client: Client) -> None: | |
| client.reset() | |
| # Create two databases with same name in different tenants | |
| admin_client = AdminClient.from_system(client._system) | |
| admin_client.create_tenant("test_tenant1") | |
| admin_client.create_tenant("test_tenant2") | |
| admin_client.create_database("test_db", tenant="test_tenant1") | |
| admin_client.create_database("test_db", tenant="test_tenant2") | |
| # Create collections in each database with same name | |
| client.set_tenant(tenant="test_tenant1", database="test_db") | |
| coll_tenant1 = client.create_collection("collection") | |
| client.set_tenant(tenant="test_tenant2", database="test_db") | |
| coll_tenant2 = client.create_collection("collection") | |
| records_tenant1 = { | |
| "ids": ["a", "b", "c"], | |
| "embeddings": [[1.0, 2.0, 3.0] for _ in range(3)], | |
| "documents": ["a", "b", "c"], | |
| } | |
| records_tenant2 = { | |
| "ids": ["c", "d", "e"], | |
| "embeddings": [[4.0, 5.0, 6.0] for _ in range(3)], | |
| "documents": ["c", "d", "e"], | |
| } | |
| # Add to the tenant1 coll | |
| coll_tenant1.add(**records_tenant1) # type: ignore | |
| # Add to the tenant2 coll | |
| coll_tenant2.add(**records_tenant2) # type: ignore | |
| # Make sure the collections are isolated | |
| res = coll_tenant1.get(include=["embeddings", "documents"]) | |
| assert res["ids"] == records_tenant1["ids"] | |
| assert res["embeddings"] == records_tenant1["embeddings"] | |
| assert res["documents"] == records_tenant1["documents"] | |
| res = coll_tenant2.get(include=["embeddings", "documents"]) | |
| assert res["ids"] == records_tenant2["ids"] | |
| assert res["embeddings"] == records_tenant2["embeddings"] | |
| assert res["documents"] == records_tenant2["documents"] | |
| def test_min_len_name(client: Client) -> None: | |
| client.reset() | |
| # Create a new database in the default tenant with a name of length 1 | |
| # and expect an error | |
| admin_client = AdminClient.from_system(client._system) | |
| with pytest.raises(Exception): | |
| admin_client.create_database("a") | |
| # Create a tenant with a name of length 1 and expect an error | |
| with pytest.raises(Exception): | |
| admin_client.create_tenant("a") | |