From bc8804b85248710b0b3cb353f85ae9d3bed3a010 Mon Sep 17 00:00:00 2001 From: Trayan Azarov <trayan.azarov@amikos.tech> Date: Fri, 27 Oct 2023 01:39:18 +0300 Subject: [PATCH] fix: Fixed an issue with test --- chromadb/test/auth/test_simple_rbac_authz.py | 53 +++++++++---------- .../basic_functionality/authz/authz.ipynb | 14 ----- 2 files changed, 25 insertions(+), 42 deletions(-) diff --git a/chromadb/test/auth/test_simple_rbac_authz.py b/chromadb/test/auth/test_simple_rbac_authz.py index 9c28024f84d..8f01b255123 100644 --- a/chromadb/test/auth/test_simple_rbac_authz.py +++ b/chromadb/test/auth/test_simple_rbac_authz.py @@ -8,9 +8,9 @@ from hypothesis import given, settings from chromadb import AdminClient -from chromadb.api import ServerAPI +from chromadb.api import AdminAPI, ServerAPI from chromadb.api.models.Collection import Collection -from chromadb.config import DEFAULT_TENANT, Settings, System +from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings, System from chromadb.test.conftest import _fastapi_fixture @@ -181,22 +181,17 @@ def token_config(draw: st.DrawFn) -> Dict[str, Any]: "db:create_database": lambda api, mapi, aapi: ( aapi.create_database(f"test-{uuid.uuid4()}") ), - "db:get_database": lambda api, mapi, aapi: ( - aapi.create_database(f"test-db-{uuid.uuid4()}"), # pre-condition - api.get_database(f"test-db-{uuid.uuid4()}"), - ), - "tenant:create_tenant": lambda api, mapi, aapi: aapi.create_tenant( - f"test-{uuid.uuid4()}" - ), - "tenant:get_tenant": lambda api, mapi, aapi: ( - aapi.create_tenant(f"test_tenant_{uuid.uuid4()}"), # pre-condition - api.get_tenant(f"test_tenant_{uuid.uuid4()}"), + "db:get_database": lambda api, mapi, aapi: (aapi.get_database(DEFAULT_DATABASE),), + "tenant:create_tenant": lambda api, mapi, aapi: ( + aapi.create_tenant(f"test-{uuid.uuid4()}") ), + "tenant:get_tenant": lambda api, mapi, aapi: (aapi.get_tenant(DEFAULT_TENANT),), "db:reset": lambda api, mapi, _: api.reset(), - "db:list_collections": lambda api, mapi: api.list_collections(), + "db:list_collections": lambda api, mapi, _: api.list_collections(), "collection:get_collection": lambda api, mapi, _: ( - mapi.create_collection(f"test-get-{uuid.uuid4()}"), # pre-condition - api.get_collection(f"test-get-{uuid.uuid4()}"), + # pre-condition + mcol := mapi.create_collection(f"test-get-{uuid.uuid4()}"), + api.get_collection(f"{mcol.name}"), ), "db:create_collection": lambda api, mapi, _: ( api.create_collection(f"test-create-{uuid.uuid4()}"), @@ -211,68 +206,68 @@ def token_config(draw: st.DrawFn) -> Dict[str, Any]: "collection:update_collection": lambda api, mapi, _: ( # pre-condition mcol := mapi.create_collection(f"test-modify-col-{uuid.uuid4()}"), - col := Collection(api, f"test-modify-col-{uuid.uuid4()}", mcol.id), + col := Collection(api, f"{mcol.name}", mcol.id), col.modify(metadata={"test": "test"}), ), "collection:add": lambda api, mapi, _: ( mcol := mapi.create_collection(f"test-add-doc-{uuid.uuid4()}"), - col := Collection(api, f"test-add-doc-{uuid.uuid4()}", mcol.id), + col := Collection(api, f"{mcol.name}", mcol.id), col.add(documents=["test"], ids=["1"]), ), "collection:delete": lambda api, mapi, _: ( mcol := mapi.create_collection(f"test-delete-doc-{uuid.uuid4()}"), mcol.add(documents=["test"], ids=["1"]), - col := Collection(api, f"test-delete-doc-{uuid.uuid4()}", mcol.id), + col := Collection(client=api, name=f"{mcol.name}", id=mcol.id), col.delete(ids=["1"]), ), "collection:get": lambda api, mapi, _: ( mcol := mapi.create_collection(f"test-get-doc-{uuid.uuid4()}"), mcol.add(documents=["test"], ids=["1"]), - col := Collection(api, f"test-get-doc-{uuid.uuid4()}", mcol.id), + col := Collection(api, f"{mcol.name}", mcol.id), col.get(ids=["1"]), ), "collection:query": lambda api, mapi, _: ( mcol := mapi.create_collection(f"test-query-doc-{uuid.uuid4()}"), mcol.add(documents=["test"], ids=["1"]), - col := Collection(api, f"test-query-doc-{uuid.uuid4()}", mcol.id), + col := Collection(api, f"{mcol.name}", mcol.id), col.query(query_texts=["test"]), ), "collection:peek": lambda api, mapi, _: ( mcol := mapi.create_collection(f"test-peek-{uuid.uuid4()}"), mcol.add(documents=["test"], ids=["1"]), - col := Collection(api, f"test-peek-{uuid.uuid4()}", mcol.id), + col := Collection(api, f"{mcol.name}", mcol.id), col.peek(), ), "collection:update": lambda api, mapi, _: ( mcol := mapi.create_collection(f"test-update-{uuid.uuid4()}"), mcol.add(documents=["test"], ids=["1"]), - col := Collection(api, f"test-update-{uuid.uuid4()}", mcol.id), + col := Collection(api, f"{mcol.name}", mcol.id), col.update(ids=["1"], documents=["test1"]), ), "collection:upsert": lambda api, mapi, _: ( mcol := mapi.create_collection(f"test-upsert-{uuid.uuid4()}"), mcol.add(documents=["test"], ids=["1"]), - col := Collection(api, f"test-upsert-{uuid.uuid4()}", mcol.id), + col := Collection(api, f"{mcol.name}", mcol.id), col.upsert(ids=["1"], documents=["test1"]), ), "collection:count": lambda api, mapi, _: ( mcol := mapi.create_collection(f"test-count-{uuid.uuid4()}"), mcol.add(documents=["test"], ids=["1"]), - col := Collection(api, f"test-count-{uuid.uuid4()}", mcol.id), + col := Collection(api, f"{mcol.name}", mcol.id), col.count(), ), } -def master_api(settings): - system = System(settings) +def master_api(_settings: Settings) -> Tuple[ServerAPI, AdminAPI]: + system = System(_settings) api = system.instance(ServerAPI) admin_api = AdminClient(api.get_settings()) system.start() return api, admin_api -@settings(max_examples=1) +@settings(max_examples=10) @given(token_config=token_config(), rbac_config=rbac_config()) def test_authz(token_config: Dict[str, Any], rbac_config: Dict[str, Any]) -> None: authz_config = rbac_config @@ -317,5 +312,7 @@ def test_authz(token_config: Dict[str, Any], rbac_config: Dict[str, Any]) -> Non "unauthorized_actions" ]: with pytest.raises(Exception) as ex: - api_executors[unauthorized_action](_api, _master_api) # type: ignore + api_executors[unauthorized_action]( + _api, _master_api, admin_api + ) # type: ignore assert "Unauthorized" in str(ex) or "Forbidden" in str(ex) diff --git a/examples/basic_functionality/authz/authz.ipynb b/examples/basic_functionality/authz/authz.ipynb index 7eab0aad970..c70df77c702 100644 --- a/examples/basic_functionality/authz/authz.ipynb +++ b/examples/basic_functionality/authz/authz.ipynb @@ -42,19 +42,6 @@ "collection.get(ids=[\"1\"])" ] }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": {}, - "outputs": [], - "source": [ - "from chromadb import AdminClient\n", - "\n", - "\n", - "admin_client = AdminClient(client.get_settings())\n", - "admin_client.create_tenant(\"qwewqewqeqwewqewqeqwewq\")" - ] - }, { "cell_type": "code", "execution_count": 25, @@ -79,7 +66,6 @@ "import uuid\n", "from chromadb.api.models.Collection import Collection\n", "\n", - "\n", "col = Collection(client, f\"test-upsert-{uuid.uuid4()}\", uuid.uuid4())\n", "col.upsert(documents=[\"test\"],ids=[\"1\"])" ]