From c4fe6ebae330f54e7ef400cf8e24650366e3688b Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Wed, 5 Feb 2025 19:10:36 +0100
Subject: [PATCH] Initial test setup.

---
 .env                                          | 12 +++++++-
 .gitignore                                    |  1 +
 run.sh                                        | 29 +++++++++++++++++++
 tests/Dockerfile.test                         |  8 +++++
 tests/open_ai/test_chat_without_tool_calls.py | 22 ++++++++++++++
 tests/pytest.ini                              |  2 ++
 tests/requirements.txt                        |  4 +++
 tests/util.py                                 | 29 +++++++++++++++++++
 8 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 tests/Dockerfile.test
 create mode 100644 tests/open_ai/test_chat_without_tool_calls.py
 create mode 100644 tests/pytest.ini
 create mode 100644 tests/requirements.txt
 create mode 100644 tests/util.py

diff --git a/.env b/.env
index 11b953d..6aaa6ca 100644
--- a/.env
+++ b/.env
@@ -1 +1,11 @@
-INVARIANT_API_URL=https://explorer.invariantlabs.ai
\ No newline at end of file
+INVARIANT_API_URL=https://explorer.invariantlabs.ai
+
+# For the testing stack only
+# TODO: Check if these can be removed.
+# Currently we rely on the `invariant explorer` command
+# to setup a local instance of explorer to test the proxy.
+# That requires these variables to be set in the .env file.
+POSTGRES_USER=postgres
+POSTGRES_PASSWORD=postgres
+POSTGRES_DB=invariantmonitor
+POSTGRES_HOST=database
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index da00cf7..3f5d309 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,6 +2,7 @@
 __pycache__/
 .pytest_cache/
 .py[oc]
+data/
 
 # Coverage and build artifacts
 .coverage
diff --git a/run.sh b/run.sh
index 1f50d77..ce3dfc3 100755
--- a/run.sh
+++ b/run.sh
@@ -20,6 +20,27 @@ down() {
   docker compose -f docker-compose.local.yml down
 }
 
+
+tests() {
+    # Run tests
+    pip install invariant-ai
+    invariant explorer up -d --build
+
+    until curl -X GET -I http://127.0.0.1/api/v1 --fail --silent --output /dev/null; do
+      echo "Backend API not available yet - checking health..."
+      sleep 2
+    done
+
+    echo "Backend API is available. Running tests..."
+
+    docker build -t 'explorer-proxy-test' -f ./tests/Dockerfile.test ./tests
+
+    docker run \
+    --mount type=bind,source=./tests,target=/tests \
+    --network host \
+    explorer-proxy-test $@
+}
+
 # -----------------------------
 # Command dispatcher
 # -----------------------------
@@ -36,4 +57,12 @@ case "$1" in
   "logs")
     docker compose -f docker-compose.local.yml logs -f
     ;;
+  "tests")
+    shift
+    tests $@
+    ;;
+  *)
+    echo "Usage: $0 [up|build|down|logs|tests]"
+    exit 1
+    ;;
 esac
\ No newline at end of file
diff --git a/tests/Dockerfile.test b/tests/Dockerfile.test
new file mode 100644
index 0000000..54476b4
--- /dev/null
+++ b/tests/Dockerfile.test
@@ -0,0 +1,8 @@
+FROM mcr.microsoft.com/playwright/python:v1.50.0-noble
+
+RUN mkdir -p /tests
+COPY ./requirements.txt /tests/requirements.txt
+WORKDIR /tests
+RUN pip install --no-cache-dir -r requirements.txt
+
+ENTRYPOINT ["pytest", "-s", "-vv"]
\ No newline at end of file
diff --git a/tests/open_ai/test_chat_without_tool_calls.py b/tests/open_ai/test_chat_without_tool_calls.py
new file mode 100644
index 0000000..5f6c7b1
--- /dev/null
+++ b/tests/open_ai/test_chat_without_tool_calls.py
@@ -0,0 +1,22 @@
+"""Test the chat completions proxy calls without tool calling."""
+
+import os
+
+# add tests folder (parent) to sys.path
+import sys
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from util import *  # needed for pytest fixtures
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+async def test_hello_world(context, url):
+    """Demo test"""
+    response = await context.request.get(
+        f"{url}/api/v1/dataset/byuser/developer/Welcome-to-Explorer"
+    )
+    dataset = await response.json()
+    assert dataset["name"] == "Welcome-to-Explorer"
+    assert dataset["user"]["username"] == "developer"
diff --git a/tests/pytest.ini b/tests/pytest.ini
new file mode 100644
index 0000000..d280de0
--- /dev/null
+++ b/tests/pytest.ini
@@ -0,0 +1,2 @@
+[pytest]
+asyncio_mode = auto
\ No newline at end of file
diff --git a/tests/requirements.txt b/tests/requirements.txt
new file mode 100644
index 0000000..9160ba9
--- /dev/null
+++ b/tests/requirements.txt
@@ -0,0 +1,4 @@
+openai
+pytest
+pytest-asyncio
+pytest-playwright
\ No newline at end of file
diff --git a/tests/util.py b/tests/util.py
new file mode 100644
index 0000000..d8bfb23
--- /dev/null
+++ b/tests/util.py
@@ -0,0 +1,29 @@
+"""Util functions for tests"""
+
+import pytest
+from playwright.async_api import async_playwright
+
+
+@pytest.fixture
+def url():
+    return "http://127.0.0.1"
+
+
+@pytest.fixture
+async def playwright(scope="session"):
+    async with async_playwright() as playwright_instance:
+        yield playwright_instance
+
+
+@pytest.fixture
+async def browser(playwright, scope="session"):
+    browser = await playwright.firefox.launch(headless=True)
+    yield browser
+    await browser.close()
+
+
+@pytest.fixture
+async def context(browser):
+    context = await browser.new_context(ignore_https_errors=True)
+    yield context
+    await context.close()