python-sdk/tests/conftest.py

from __future__ import annotations

import os
import sys
import logging
from typing import TYPE_CHECKING, Iterator, AsyncIterator

# Add pkg directory to path for hanzoai import
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__)), "pkg"))

import pytest
import pytest_asyncio
from pytest_asyncio import is_async_test

from hanzoai import Hanzo, AsyncHanzo
from tests.mock_fixtures import mock_api

if TYPE_CHECKING:
    from _pytest.fixtures import FixtureRequest

pytest.register_assert_rewrite("tests.utils")

logging.getLogger("hanzoai").setLevel(logging.DEBUG)


# automatically add `pytest.mark.asyncio()` to all of our async tests
# so we don't have to add that boilerplate everywhere
def pytest_collection_modifyitems(items: list[pytest.Function]) -> None:
    pytest_asyncio_tests = (item for item in items if is_async_test(item))
    session_scope_marker = pytest.mark.asyncio()
    for async_test in pytest_asyncio_tests:
        async_test.add_marker(session_scope_marker, append=False)


base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")

api_key = "My API Key"


@pytest.fixture(scope="function")
def client(request: FixtureRequest, mock_api) -> Iterator[Hanzo]:
    strict = getattr(request, "param", True)
    if not isinstance(strict, bool):
        raise TypeError(f"Unexpected fixture parameter type {type(strict)}, expected {bool}")

    with Hanzo(base_url=base_url, api_key=api_key, _strict_response_validation=strict) as client:
        yield client


@pytest_asyncio.fixture(scope="function")
async def async_client(request: FixtureRequest, mock_api) -> AsyncIterator[AsyncHanzo]:
    strict = getattr(request, "param", True)
    if not isinstance(strict, bool):
        raise TypeError(f"Unexpected fixture parameter type {type(strict)}, expected {bool}")

    async with AsyncHanzo(base_url=base_url, api_key=api_key, _strict_response_validation=strict) as client:
        yield client