Skip to content

Commit

Permalink
fix: Fixed an issue with test
Browse files Browse the repository at this point in the history
  • Loading branch information
@tazarov
tazarov committed Oct 26, 2023
1 parent 38fd73a commit bc8804b
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 42 deletions.
53 changes: 25 additions & 28 deletions chromadb/test/auth/test_simple_rbac_authz.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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()}"),
Expand All @@ -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
Expand Down Expand Up @@ -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)
14 changes: 0 additions & 14 deletions examples/basic_functionality/authz/authz.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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\"])"
]
Expand Down

0 comments on commit bc8804b

Please sign in to comment.