implement tests for being able to create collections

cli/get_token.py

@@ -1,12 +1,8 @@
-from concierge_backend_lib.authentication import (
-    get_keycloak_client,
-    get_keycloak_admin_openid_token,
-)
+from concierge_backend_lib.authentication import get_keycloak_admin_openid_token
 from concierge_backend_lib.authorization import auth_enabled
 
 
 def get_token():
     if not auth_enabled:
         return {"access_token": None}
-    openid_client = get_keycloak_client()
-    return get_keycloak_admin_openid_token(openid_client)
+    return get_keycloak_admin_openid_token()

concierge_backend_lib/authentication.py

@@ -67,7 +67,8 @@ def get_keycloak_admin_client():
     return client
 
 
-def get_keycloak_admin_openid_token(client: KeycloakOpenID):
+def get_keycloak_admin_openid_token():
+    client = get_keycloak_client()
     token = client.token(grant_type="client_credentials")
     return token
 

concierge_backend_lib/authorization.py

@@ -8,7 +8,7 @@
 import os
 from keycloak import KeycloakPostError
 from keycloak.exceptions import raise_error_from_response
-from .httpx_client import httpx_client
+import httpx
 
 
 class UnauthorizedOperationError(Exception):
@@ -24,16 +24,19 @@ async def authorize(token, resource, scope: str | None = None):
     permission = resource
     if scope:
         permission += f"#{scope}"
-    resp = await httpx_client.post(
-        keycloak_openid_config()["token_endpoint"],
-        data={
-            "grant_type": "urn:ietf:params:oauth:grant-type:uma-ticket",
-            "audience": "concierge-auth",
-            "permission": permission,
-            "response_mode": "decision",
-        },
-        headers={"Authorization": f"Bearer {token}"},
-    )
+    async with httpx.AsyncClient(
+        verify=os.getenv("ROOT_CA"), timeout=None
+    ) as httpx_client:
+        resp = await httpx_client.post(
+            keycloak_openid_config()["token_endpoint"],
+            data={
+                "grant_type": "urn:ietf:params:oauth:grant-type:uma-ticket",
+                "audience": "concierge-auth",
+                "permission": permission,
+                "response_mode": "decision",
+            },
+            headers={"Authorization": f"Bearer {token}"},
+        )
 
     # this will cause us to get a new token if needed
     raise_error_from_response(resp, KeycloakPostError)

concierge_backend_lib/prompting.py

@@ -5,8 +5,7 @@
 from .opensearch_prompting import get_context_from_opensearch
 from .authorization import auth_enabled, authorize, UnauthorizedOperationError
 from isi_util.async_single import asyncify
-from .httpx_client import httpx_client
-import traceback
+import httpx
 
 load_dotenv()
 HOST = os.getenv("OLLAMA_HOST") or "localhost"
@@ -66,13 +65,12 @@ async def get_response(
 
     data = {"model": "mistral", "prompt": prompt, "stream": False}
     # Ollama doesn't like the data in JSON, we have to dump it to string
-    try:
+    async with httpx.AsyncClient(
+        verify=os.getenv("ROOT_CA"), timeout=None
+    ) as httpx_client:
         response = await httpx_client.post(
             f"http://{HOST}:11434/api/generate", data=json.dumps(data)
         )
-    except Exception:
-        traceback.print_exc()
-
     if response.status_code != 200:
         print(response.content)
         return f"ollama status: {response.status_code}"

dev_requirements.txt

@@ -4,6 +4,7 @@ pipdeptree~=2.23.3
 twine~=5.1.1
 zxcvbn~=4.4.28
 pytest
+pytest-asyncio
 -e concierge_packages/script_builder
 -e concierge_packages/installer
 -r local_requirements.txt

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+asyncio_mode=auto
+asyncio_default_fixture_loop_scope=session

tests/test_rbac.py

@@ -0,0 +1,70 @@
+# This test should be run once Concierge is installed with the demo RBAC configuration
+# For best results it should be fresh installation with no collections created or tweaks made to the access controls
+# Do not use this on a production instance!
+
+from concierge_backend_lib.authentication import (
+    get_keycloak_client,
+    get_keycloak_admin_openid_token,
+)
+from concierge_backend_lib.document_collections import (
+    create_collection,
+    delete_collection,
+)
+import pytest
+from keycloak import KeycloakPostError
+import asyncio
+
+keycloak_client = get_keycloak_client()
+
+# we will use the collections created in the first tests for the subsequent tests
+collection_lookup = {}
+
+
+async def create_collection_for_user(user, location):
+    token = keycloak_client.token(user, "test")
+    collection_name = f"{user}'s {location} collection"
+    collection_id = await create_collection(
+        token["access_token"], f"{user}'s {location} collection", location
+    )
+    collection_lookup[collection_name] = collection_id
+    return collection_id
+
+
+@pytest.mark.parametrize(
+    "user,location",
+    [
+        ("testadmin", "private"),
+        ("testadmin", "shared"),
+        ("testshared", "shared"),
+        ("testprivate", "private"),
+    ],
+)
+async def test_can_create_collection(user, location):
+    assert await create_collection_for_user(user, location)
+
+
+@pytest.mark.parametrize(
+    "user,location",
+    [
+        ("testshared", "private"),
+        ("testprivate", "shared"),
+        ("testsharedread", "private"),
+        ("testsharedread", "shared"),
+        ("testnothing", "private"),
+        ("testnothing", "shared"),
+    ],
+)
+async def test_cannot_create_collection(user, location):
+    with pytest.raises(KeycloakPostError):
+        await create_collection_for_user(user, location)
+
+
+async def teardown():
+    token = get_keycloak_admin_openid_token()
+    for id in collection_lookup.values():
+        token = get_keycloak_admin_openid_token()
+        await delete_collection(token["access_token"], id)
+
+
+def teardown_module():
+    asyncio.run(teardown())