feat: add comprehensive test suite for BrainusAI client and models

Author: 0xneobyte

.github/workflows/tests.yml

@@ -0,0 +1,29 @@
+name: Tests
+
+on:
+  push:
+    branches: ["main", "dev"]
+  pull_request:
+    branches: ["main", "dev"]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: pip install -e ".[dev]"
+
+      - name: Run tests
+        run: python -m pytest tests/ -v

pyproject.toml

@@ -60,6 +60,10 @@ target-version = "py311"
 select = ["E", "F", "I", "UP", "B", "SIM", "TCH"]
 ignore = []
 
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]
+
 [tool.mypy]
 python_version = "3.11"
 strict = true

src/brainus_ai/__init__.py

@@ -14,7 +14,7 @@
 )
 from .models import QueryRequest, QueryResponse, Citation, UsageStats, Plan, QueryFilters, PlanInfo
 
-__version__ = "0.1.0"
+__version__ = "0.1.6"
 
 __all__ = [
     "BrainusAI",

tests/__init__.py

@@ -0,0 +1,79 @@
+"""Shared fixtures for Brainus AI SDK tests."""
+
+import pytest
+from brainus_ai import BrainusAI
+
+VALID_API_KEY = "brainus_test_key_abc123"
+BASE_URL = "https://api.brainus.lk"
+
+
+@pytest.fixture
+def api_key() -> str:
+    return VALID_API_KEY
+
+
+@pytest.fixture
+def client() -> BrainusAI:
+    return BrainusAI(api_key=VALID_API_KEY)
+
+
+# Reusable mock response payloads
+QUERY_RESPONSE = {
+    "answer": "Python is a high-level programming language.",
+    "citations": [
+        {
+            "document_id": "doc_001",
+            "document_name": "Python Basics.pdf",
+            "pages": [1, 2],
+            "metadata": {"subject": "ICT", "grade": "12"},
+            "chunk_text": "Python is a high-level...",
+        }
+    ],
+    "has_citations": True,
+}
+
+USAGE_RESPONSE = {
+    "total_requests": 42,
+    "total_tokens": 8500,
+    "total_cost_usd": 0.034,
+    "by_endpoint": {"/api/v1/dev/query": 42},
+    "quota_remaining": 958,
+    "quota_percentage": 4.2,
+    "plan": {
+        "name": "Pro",
+        "rate_limit_per_minute": 60,
+        "rate_limit_per_day": 1000,
+        "monthly_quota": 1000,
+    },
+    "period_start": "2026-03-01",
+    "period_end": "2026-03-31",
+}
+
+PLANS_RESPONSE = {
+    "plans": [
+        {
+            "id": "plan_free",
+            "name": "Free",
+            "description": "Free tier",
+            "rate_limit_per_minute": 10,
+            "rate_limit_per_day": 100,
+            "monthly_quota": 100,
+            "price_lkr": None,
+            "allowed_models": ["brainusai-fast"],
+            "features": {},
+            "is_active": True,
+        },
+        {
+            "id": "plan_pro",
+            "name": "Pro",
+            "description": "Pro tier",
+            "rate_limit_per_minute": 60,
+            "rate_limit_per_day": 1000,
+            "monthly_quota": 1000,
+            "price_lkr": 2500.0,
+            "allowed_models": ["brainusai-fast", "brainusai-thinking"],
+            "features": {"priority_support": True},
+            "is_active": True,
+        },
+    ]
+}

tests/conftest.py

@@ -0,0 +1,55 @@
+"""Tests for BrainusAI client initialization."""
+
+import pytest
+from brainus_ai import BrainusAI
+from brainus_ai.exceptions import AuthenticationError
+
+
+class TestClientInit:
+    def test_valid_key_accepted(self):
+        client = BrainusAI(api_key="brainus_abc123")
+        assert client.api_key == "brainus_abc123"
+
+    def test_empty_key_raises(self):
+        with pytest.raises(AuthenticationError):
+            BrainusAI(api_key="")
+
+    def test_wrong_prefix_sk_live_raises(self):
+        with pytest.raises(AuthenticationError):
+            BrainusAI(api_key="sk_live_abc123")
+
+    def test_wrong_prefix_bare_string_raises(self):
+        with pytest.raises(AuthenticationError):
+            BrainusAI(api_key="my_secret_key")
+
+    def test_wrong_prefix_openai_style_raises(self):
+        with pytest.raises(AuthenticationError):
+            BrainusAI(api_key="sk-proj-abc123")
+
+    def test_default_base_url(self):
+        client = BrainusAI(api_key="brainus_abc123")
+        assert client.base_url == "https://api.brainus.lk"
+
+    def test_trailing_slash_stripped_from_base_url(self):
+        client = BrainusAI(api_key="brainus_abc123", base_url="https://api.brainus.lk/")
+        assert client.base_url == "https://api.brainus.lk"
+
+    def test_custom_base_url(self):
+        client = BrainusAI(api_key="brainus_abc123", base_url="http://localhost:8000")
+        assert client.base_url == "http://localhost:8000"
+
+    def test_default_timeout(self):
+        client = BrainusAI(api_key="brainus_abc123")
+        assert client.timeout == 30.0
+
+    def test_custom_timeout(self):
+        client = BrainusAI(api_key="brainus_abc123", timeout=60.0)
+        assert client.timeout == 60.0
+
+    def test_default_max_retries(self):
+        client = BrainusAI(api_key="brainus_abc123")
+        assert client.max_retries == 3
+
+    def test_custom_max_retries(self):
+        client = BrainusAI(api_key="brainus_abc123", max_retries=5)
+        assert client.max_retries == 5

tests/test_client.py

@@ -0,0 +1,39 @@
+"""Tests for async context manager and close()."""
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+from pytest_httpx import HTTPXMock
+
+from brainus_ai import BrainusAI
+
+from .conftest import BASE_URL, QUERY_RESPONSE, VALID_API_KEY
+
+QUERY_URL = f"{BASE_URL}/api/v1/dev/query"
+
+
+async def test_context_manager_returns_client() -> None:
+    async with BrainusAI(api_key=VALID_API_KEY) as client:
+        assert isinstance(client, BrainusAI)
+
+
+async def test_context_manager_closes_on_exit(httpx_mock: HTTPXMock) -> None:
+    httpx_mock.add_response(method="GET", url=f"{BASE_URL}/api/v1/dev/usage", json={"total_requests": 0, "by_endpoint": {}})
+    async with BrainusAI(api_key=VALID_API_KEY) as client:
+        await client.get_usage()
+    # After exiting, the underlying httpx client should be closed.
+    assert client._client.is_closed
+
+
+async def test_close_marks_client_as_closed() -> None:
+    client = BrainusAI(api_key=VALID_API_KEY)
+    assert not client._client.is_closed
+    await client.close()
+    assert client._client.is_closed
+
+
+async def test_requests_work_inside_context_manager(httpx_mock: HTTPXMock) -> None:
+    httpx_mock.add_response(method="POST", url=QUERY_URL, json=QUERY_RESPONSE)
+    async with BrainusAI(api_key=VALID_API_KEY) as client:
+        resp = await client.query("test")
+    assert resp.answer == "Python is a high-level programming language."

tests/test_context_manager.py

@@ -0,0 +1,96 @@
+"""Tests for exception classes."""
+
+from brainus_ai.exceptions import (
+    APIError,
+    AuthenticationError,
+    BrainusError,
+    QuotaExceededError,
+    RateLimitError,
+)
+
+
+class TestBrainusError:
+    def test_is_exception(self) -> None:
+        assert issubclass(BrainusError, Exception)
+
+    def test_message_attribute(self) -> None:
+        e = BrainusError("something broke")
+        assert e.message == "something broke"
+
+    def test_str_representation(self) -> None:
+        e = BrainusError("something broke")
+        assert "something broke" in str(e)
+
+
+class TestAuthenticationError:
+    def test_is_brainus_error(self) -> None:
+        assert issubclass(AuthenticationError, BrainusError)
+
+    def test_message(self) -> None:
+        e = AuthenticationError("bad key")
+        assert e.message == "bad key"
+
+
+class TestRateLimitError:
+    def test_is_brainus_error(self) -> None:
+        assert issubclass(RateLimitError, BrainusError)
+
+    def test_retry_after_none_by_default(self) -> None:
+        e = RateLimitError()
+        assert e.retry_after is None
+
+    def test_retry_after_set(self) -> None:
+        e = RateLimitError("limit hit", retry_after=60)
+        assert e.retry_after == 60
+
+    def test_default_message(self) -> None:
+        e = RateLimitError()
+        assert e.message == "Rate limit exceeded"
+
+    def test_custom_message(self) -> None:
+        e = RateLimitError("slow down")
+        assert e.message == "slow down"
+
+
+class TestQuotaExceededError:
+    def test_is_brainus_error(self) -> None:
+        assert issubclass(QuotaExceededError, BrainusError)
+
+    def test_message(self) -> None:
+        e = QuotaExceededError("quota gone")
+        assert e.message == "quota gone"
+
+
+class TestAPIError:
+    def test_is_brainus_error(self) -> None:
+        assert issubclass(APIError, BrainusError)
+
+    def test_status_code_none_by_default(self) -> None:
+        e = APIError("error")
+        assert e.status_code is None
+
+    def test_status_code_set(self) -> None:
+        e = APIError("server error", status_code=500)
+        assert e.status_code == 500
+
+    def test_message(self) -> None:
+        e = APIError("bad gateway", status_code=502)
+        assert e.message == "bad gateway"
+
+
+class TestInstanceofChecks:
+    def test_all_subclass_of_brainus_error(self) -> None:
+        for cls in [AuthenticationError, RateLimitError, QuotaExceededError, APIError]:
+            assert issubclass(cls, BrainusError)
+
+    def test_all_subclass_of_exception(self) -> None:
+        for cls in [BrainusError, AuthenticationError, RateLimitError, QuotaExceededError, APIError]:
+            assert issubclass(cls, Exception)
+
+    def test_can_catch_with_base_class(self) -> None:
+        try:
+            raise AuthenticationError("bad key")
+        except BrainusError as e:
+            assert e.message == "bad key"
+        else:
+            assert False, "Should have caught exception"