fix: fixed integration tests (#147)

.env.example

@@ -7,5 +7,4 @@ DIAL_URL=<dial core url>
 
 # Misc env vars for the server
 LOG_LEVEL=INFO # Default in prod is INFO. Use DEBUG for dev.
-WEB_CONCURRENCY=1 # Number of uvicorn workers
-TEST_SERVER_URL=http://0.0.0.0:5001
+WEB_CONCURRENCY=1 # Number of uvicorn workers

README.md

@@ -12,10 +12,10 @@ Note that a model supports `/truncate_prompt` endpoint if and only if it support
 
 |Vendor|Model|Deployment name|Modality|`/tokenize`|`/truncate_prompt`, `max_prompt_tokens`|tools/functions|
 |---|---|---|---|---|---|---|
-|Anthropic|Claude 3.5 Sonnet|anthropic.claude-3-5-sonnet-20240620-v1:0|text-to-text, image-to-text|🟡|🟡|✅|
-|Anthropic|Claude 3 Sonnet|anthropic.claude-3-sonnet-20240229-v1:0|text-to-text, image-to-text|🟡|🟡|✅|
-|Anthropic|Claude 3 Haiku|anthropic.claude-3-haiku-20240307-v1:0|text-to-text, image-to-text|🟡|🟡|✅|
-|Anthropic|Claude 3 Opus|anthropic.claude-3-opus-20240229-v1:0|text-to-text, image-to-text|🟡|🟡|✅|
+|Anthropic|Claude 3.5 Sonnet|[us.\|eu.]anthropic.claude-3-5-sonnet-20240620-v1:0|text-to-text, image-to-text|🟡|🟡|✅|
+|Anthropic|Claude 3 Sonnet|[us.\|eu.]anthropic.claude-3-sonnet-20240229-v1:0|text-to-text, image-to-text|🟡|🟡|✅|
+|Anthropic|Claude 3 Haiku|[us.\|eu.]anthropic.claude-3-haiku-20240307-v1:0|text-to-text, image-to-text|🟡|🟡|✅|
+|Anthropic|Claude 3 Opus|[us.]anthropic.claude-3-opus-20240229-v1:0|text-to-text, image-to-text|🟡|🟡|✅|
 |Anthropic|Claude 2.1|anthropic.claude-v2:1|text-to-text|✅|✅|✅|
 |Anthropic|Claude 2|anthropic.claude-v2|text-to-text|✅|✅|❌|
 |Anthropic|Claude Instant 1.2|anthropic.claude-instant-v1|text-to-text|🟡|🟡|❌|
@@ -105,7 +105,6 @@ Copy `.env.example` to `.env` and customize it for your environment:
 |AIDIAL_LOG_LEVEL|WARNING|AI DIAL SDK log level|
 |DIAL_URL||URL of the core DIAL server. If defined, images generated by Stability are uploaded to the DIAL file storage and attachments are returned with URLs pointing to the images. Otherwise, the images are returned as base64 encoded strings.|
 |WEB_CONCURRENCY|1|Number of workers for the server|
-|TEST_SERVER_URL|http://0.0.0.0:5001|Server URL used in the integration tests|
 
 ## Load balancing
 

aidial_adapter_bedrock/aws_client_config.py

@@ -2,7 +2,7 @@
 
 import boto3
 from aidial_sdk.embeddings import Request
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from aidial_adapter_bedrock.utils.concurrency import make_async
 from aidial_adapter_bedrock.utils.env import get_aws_default_region
@@ -43,7 +43,7 @@ def get_anthropic_bedrock_client_kwargs(self) -> dict:
 
 
 class UpstreamConfig(BaseModel):
-    region: str = get_aws_default_region()
+    region: str = Field(default_factory=get_aws_default_region)
     aws_access_key_id: str | None = None
     aws_secret_access_key: str | None = None
     aws_assume_role_arn: str | None = os.environ.get("AWS_ASSUME_ROLE_ARN")

aidial_adapter_bedrock/llm/truncate_prompt.py

@@ -51,9 +51,8 @@ def to_dial_exception(self) -> DialException:
 
 
 def _partition_indexer(chunks: List[int]) -> Callable[[int], List[int]]:
-    """Returns a function that maps an index to indices of its partition.
-    >>> [_partition_indexer([2, 3])(i) for i in range(5)]
-    [[0, 1], [0, 1], [2, 3, 4], [2, 3, 4], [2, 3, 4]]
+    """
+    Returns a function that maps an index to indices of its partition.
     """
     mapping: dict[int, List[int]] = {}
     offset = 0

aidial_adapter_bedrock/utils/env.py

@@ -1,4 +1,5 @@
 import os
+from functools import cache
 from typing import Optional
 
 from aidial_adapter_bedrock.utils.log_config import app_logger as log
@@ -13,6 +14,7 @@ def get_env(name: str, err_msg: Optional[str] = None) -> str:
     raise Exception(err_msg or f"{name} env variable is not set")
 
 
+@cache
 def get_aws_default_region() -> str:
     region = os.getenv("DEFAULT_REGION")
     if region is not None:

noxfile.py

@@ -35,7 +35,7 @@ def format(session: nox.Session):
 
 def run_tests(session: nox.Session, *args):
     session.run("poetry", "install", external=True)
-    session.run("pytest", "aidial_adapter_bedrock", *args)
+    session.run("pytest", *args)
 
 
 @nox.session
@@ -45,4 +45,4 @@ def test(session: nox.Session):
 
 @nox.session
 def integration_tests(session: nox.Session):
-    run_tests(session, "-n=auto", "tests/integration_tests/")
+    run_tests(session, "tests/integration_tests/")

pyproject.toml

@@ -45,7 +45,8 @@ flake8 = "6.0.0"
 nox = "^2023.4.22"
 
 [tool.pytest.ini_options]
-addopts = "--doctest-modules"
+addopts="-n=auto"
+env_override_existing_values = 1
 # muting warnings coming from opentelemetry package
 filterwarnings = [
     "ignore::DeprecationWarning:opentelemetry.instrumentation.dependencies"

tests/conftest.py

@@ -1,15 +1,47 @@
-import os
+from typing import Mapping
 
+import httpx
 import pytest
+import pytest_asyncio
+from httpx import ASGITransport
+from openai import AsyncAzureOpenAI
 
-from tests.utils.server import server_generator
 
-DEFAULT_API_VERSION = "2023-03-15-preview"
-TEST_SERVER_URL = os.getenv("TEST_SERVER_URL", "http://0.0.0.0:5001")
+@pytest.fixture(autouse=True)
+def configure_unit_tests(monkeypatch, request):
+    """
+    Set up fake environment variables for unit tests.
+    """
+    if "tests/unit_tests" in request.node.nodeid:
+        monkeypatch.setenv("AWS_DEFAULT_REGION", "test-region")
 
 
-@pytest.fixture(scope="module")
-def server():
-    yield from server_generator(
-        "aidial_adapter_bedrock.app:app", TEST_SERVER_URL
-    )
+@pytest_asyncio.fixture
+async def test_http_client():
+    from aidial_adapter_bedrock.app import app
+
+    async with httpx.AsyncClient(
+        transport=ASGITransport(app),  # type: ignore
+        base_url="http://test-app.com",
+    ) as client:
+        yield client
+
+
+@pytest.fixture
+def get_openai_client(test_http_client: httpx.AsyncClient):
+    def _get_client(
+        deployment_id: str | None = None,
+        extra_headers: Mapping[str, str] | None = None,
+    ) -> AsyncAzureOpenAI:
+        return AsyncAzureOpenAI(
+            azure_endpoint=str(test_http_client.base_url),
+            azure_deployment=deployment_id,
+            api_version="",
+            api_key="dummy_key",
+            max_retries=2,
+            timeout=30,
+            http_client=test_http_client,
+            default_headers=extra_headers,
+        )
+
+    yield _get_client