102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
import pytest
|
|
|
|
from tht.adapters.dwh import PostgresDwhAdapter, ThothRestDwhAdapter
|
|
from tht.adapters.vector import LegacyDirectVectorStore, ThothHttpVectorStore
|
|
from tht.adapters.factory import build_dwh, build_vector_store
|
|
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",
|
|
"connection": {
|
|
"host": "vector-db",
|
|
"database": "postgres",
|
|
"schema": "vectors",
|
|
"user": "reader",
|
|
"password": "secret",
|
|
},
|
|
}
|
|
)
|
|
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_reader():
|
|
config = _config(vector_type="pgvector_direct")
|
|
|
|
assert isinstance(build_vector_store(config), LegacyDirectVectorStore)
|
|
with pytest.raises(ConfigError, match="writer"):
|
|
build_vector_store(config, 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
|