ThothII/harness/tests/test_adapter_factory.py

import pytest

from tht.adapters.dwh import PostgresDwhAdapter, ThothRestDwhAdapter
from tht.adapters.factory import build_dwh, build_vector_store
from tht.adapters.vector import PgVectorStore, QdrantVectorStore, ThothHttpVectorStore
from tht.config import Config, ConfigError


def _config(*, dwh_type="thoth_rest", vector_type="thoth_vector_http", reader=True, writer=True):
    dwh = (
        {
            "type": "thoth_rest",
            "database": {"database": "analytics", "schema": "mart"},
            "endpoint": {"base_url": "https://dwh.test/", "api_key": "reader"},
        }
        if dwh_type == "thoth_rest"
        else {
            "type": "postgres_direct",
            "connection": {
                "host": "db",
                "database": "analytics",
                "schema": "mart",
                "user": "reader",
                "password": "secret",
            },
        }
    )
    vectors = (
        {
            "type": "thoth_vector_http",
            **(
                {"reader": {"base_url": "https://vectors.test/", "api_key": "reader"}}
                if reader
                else {}
            ),
            **(
                {"writer": {"base_url": "https://vectors.test/", "api_key": "writer"}}
                if writer
                else {}
            ),
        }
        if vector_type == "thoth_vector_http"
        else {
            "type": "pgvector_direct",
            **(
                {
                    "reader": {
                        "host": "vector-db",
                        "database": "postgres",
                        "schema": "vectors",
                        "user": "reader",
                        "password": "secret",
                    }
                }
                if reader
                else {}
            ),
            **(
                {
                    "writer": {
                        "host": "vector-db",
                        "database": "postgres",
                        "schema": "vectors",
                        "user": "writer",
                        "password": "secret",
                    }
                }
                if writer
                else {}
            ),
        }
    )
    legacy_database = (
        dwh["connection"]
        if dwh_type == "postgres_direct"
        else {
            **dwh["database"],
            "user": "rest",
            "password": "",
            "transport": "rest",
        }
    )
    return Config.model_validate({"dwh": dwh, "vectors": vectors, "database": legacy_database})


@pytest.mark.parametrize(
    ("dwh_type", "adapter_type"),
    [("postgres_direct", PostgresDwhAdapter), ("thoth_rest", ThothRestDwhAdapter)],
)
def test_factory_selects_dwh_adapter(dwh_type, adapter_type):
    assert isinstance(build_dwh(_config(dwh_type=dwh_type)), adapter_type)


def test_factory_selects_http_vector_and_requires_writer():
    config = _config(writer=False)

    assert isinstance(build_vector_store(config), ThothHttpVectorStore)
    with pytest.raises(ConfigError, match="writer"):
        build_vector_store(config, require_write=True)


def test_factory_builds_writer_only_http_vector_when_write_is_required():
    config = _config(reader=False, writer=True)

    store = build_vector_store(config, require_write=True)
    assert isinstance(store, ThothHttpVectorStore)
    assert store.capabilities.search is False
    assert store.capabilities.upsert is True


def test_factory_selects_direct_vector_store_and_requires_writer():
    config = _config(vector_type="pgvector_direct", writer=False)

    assert isinstance(build_vector_store(config), PgVectorStore)
    with pytest.raises(ConfigError, match="writer"):
        build_vector_store(config, require_write=True)


def test_factory_builds_writer_only_direct_vector_when_write_is_required():
    store = build_vector_store(
        _config(vector_type="pgvector_direct", reader=False), require_write=True
    )
    assert isinstance(store, PgVectorStore)
    assert store.capabilities.search is False
    assert store.capabilities.upsert is True


def test_factory_selects_qdrant_for_schema_v3_runtime():
    config = Config.model_validate(
        {
            "dwh": {
                "type": "postgres_direct",
                "connection": {
                    "host": "db",
                    "database": "analytics",
                    "schema": "mart",
                    "user": "reader",
                    "password": "secret",
                },
            },
            "database": {
                "host": "db",
                "database": "analytics",
                "schema": "mart",
                "user": "reader",
                "password": "secret",
                "transport": "direct",
            },
            "vectors": {
                "type": "qdrant",
                "base_url": "http://qdrant:6333",
                "collection": "psd-clinical",
            },
            "embeddings": {
                "provider": "ollama_internal",
                "base_url": "http://embedding:11434",
                "model": "qwen3-embedding:0.6b",
                "dim": 1024,
            },
        }
    )
    config._workspace_id = "psd-clinical"
    config._workspace_revision = "a" * 40

    store = build_vector_store(config, require_write=True)

    assert isinstance(store, QdrantVectorStore)
    assert store.capabilities.search is True
    assert store.capabilities.upsert is True


def test_factory_reuses_legacy_direct_connection_for_server_writes_only():
    server = _config(vector_type="pgvector_direct", writer=False)
    server.vectors.connection = server.vectors.reader
    server.vectors.reader = None

    store = build_vector_store(server, require_write=True)
    assert store.capabilities.search is True
    assert store.capabilities.upsert is True

    server.profile = "workstation"
    with pytest.raises(ConfigError, match="writer"):
        build_vector_store(server, require_write=True)


def test_factory_propagates_non_default_statement_timeout():
    config = _config(dwh_type="postgres_direct")
    config.execution.statement_timeout_ms = 12_345
    assert build_dwh(config)._statement_timeout_ms == 12_345