From 6b1a9fec11c8c73fd37d26e6feee56905d7deaa1 Mon Sep 17 00:00:00 2001 From: sebovzeoueb <7989595+sebovzeoueb@users.noreply.github.com> Date: Mon, 16 Dec 2024 18:23:34 +0100 Subject: [PATCH] implement tests for being able to create collections --- cli/get_token.py | 8 +-- concierge_backend_lib/authentication.py | 3 +- concierge_backend_lib/authorization.py | 25 +++++---- concierge_backend_lib/httpx_client.py | 6 --- concierge_backend_lib/prompting.py | 10 ++-- dev_requirements.txt | 1 + pytest.ini | 3 ++ tests/test_rbac.py | 70 +++++++++++++++++++++++++ 8 files changed, 96 insertions(+), 30 deletions(-) delete mode 100644 concierge_backend_lib/httpx_client.py create mode 100644 pytest.ini create mode 100644 tests/test_rbac.py diff --git a/cli/get_token.py b/cli/get_token.py index 3cb4efe..3fe7015 100644 --- a/cli/get_token.py +++ b/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() diff --git a/concierge_backend_lib/authentication.py b/concierge_backend_lib/authentication.py index 0582e58..8aab677 100644 --- a/concierge_backend_lib/authentication.py +++ b/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 diff --git a/concierge_backend_lib/authorization.py b/concierge_backend_lib/authorization.py index 047018f..77cf2a2 100644 --- a/concierge_backend_lib/authorization.py +++ b/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) diff --git a/concierge_backend_lib/httpx_client.py b/concierge_backend_lib/httpx_client.py deleted file mode 100644 index 933bdc6..0000000 --- a/concierge_backend_lib/httpx_client.py +++ /dev/null @@ -1,6 +0,0 @@ -import httpx -import os -from dotenv import load_dotenv - -load_dotenv() -httpx_client = httpx.AsyncClient(verify=os.getenv("ROOT_CA"), timeout=None) diff --git a/concierge_backend_lib/prompting.py b/concierge_backend_lib/prompting.py index 051988f..bb78a6b 100644 --- a/concierge_backend_lib/prompting.py +++ b/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}" diff --git a/dev_requirements.txt b/dev_requirements.txt index b32eb2b..0386538 100644 --- a/dev_requirements.txt +++ b/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 \ No newline at end of file diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..38b81c1 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +asyncio_mode=auto +asyncio_default_fixture_loop_scope=session \ No newline at end of file diff --git a/tests/test_rbac.py b/tests/test_rbac.py new file mode 100644 index 0000000..a428f2c --- /dev/null +++ b/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())