146 lines
4.7 KiB
Python
146 lines
4.7 KiB
Python
import pytest
|
|
|
|
from tht.adapters.dwh import PostgresDwhAdapter, ThothRestDwhAdapter
|
|
from tht.adapters.vector import PgVectorStore, 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",
|
|
**(
|
|
{
|
|
"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_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
|