132 lines
3.7 KiB
Python
132 lines
3.7 KiB
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
import yaml
|
|
from pydantic import ValidationError
|
|
|
|
from haiku.rag.config import (
|
|
AppConfig,
|
|
FSSourceConfig,
|
|
HTTPSourceConfig,
|
|
IngesterConfig,
|
|
RetryPolicyConfig,
|
|
S3SourceConfig,
|
|
)
|
|
|
|
|
|
def test_default_ingester_config_has_sane_values():
|
|
cfg = IngesterConfig()
|
|
assert cfg.sources == []
|
|
assert cfg.workers.worker_count == 4
|
|
assert cfg.workers.retry.max_attempts == 5
|
|
assert cfg.api.enabled is True
|
|
assert cfg.api.port == 8765
|
|
|
|
|
|
def test_discriminator_picks_fs_source():
|
|
cfg = IngesterConfig.model_validate(
|
|
{"sources": [{"type": "fs", "root": "/tmp/docs"}]}
|
|
)
|
|
assert isinstance(cfg.sources[0], FSSourceConfig)
|
|
assert cfg.sources[0].root == Path("/tmp/docs")
|
|
assert cfg.sources[0].delete_orphans is True
|
|
|
|
|
|
def test_discriminator_picks_http_source():
|
|
cfg = IngesterConfig.model_validate(
|
|
{"sources": [{"type": "http", "id": "arxiv", "urls": ["https://x"]}]}
|
|
)
|
|
assert isinstance(cfg.sources[0], HTTPSourceConfig)
|
|
assert cfg.sources[0].id == "arxiv"
|
|
|
|
|
|
def test_discriminator_picks_s3_source():
|
|
cfg = IngesterConfig.model_validate(
|
|
{"sources": [{"type": "s3", "uri": "s3://bucket/prefix/"}]}
|
|
)
|
|
assert isinstance(cfg.sources[0], S3SourceConfig)
|
|
assert cfg.sources[0].uri == "s3://bucket/prefix/"
|
|
|
|
|
|
def test_discriminator_picks_webdav_source():
|
|
cfg = IngesterConfig.model_validate(
|
|
{
|
|
"sources": [
|
|
{
|
|
"type": "webdav",
|
|
"id": "nc",
|
|
"base_url": "https://nc.example.com/dav/",
|
|
"username": "alice",
|
|
"password": "hunter2",
|
|
}
|
|
]
|
|
}
|
|
)
|
|
from haiku.rag.config import WebDAVSourceConfig
|
|
|
|
assert isinstance(cfg.sources[0], WebDAVSourceConfig)
|
|
assert cfg.sources[0].base_url == "https://nc.example.com/dav/"
|
|
assert cfg.sources[0].username == "alice"
|
|
|
|
|
|
def test_discriminator_rejects_unknown_type():
|
|
with pytest.raises(ValidationError):
|
|
IngesterConfig.model_validate({"sources": [{"type": "ftp", "uri": "x"}]})
|
|
|
|
|
|
def test_per_source_retry_overrides_default():
|
|
cfg = IngesterConfig.model_validate(
|
|
{
|
|
"sources": [
|
|
{
|
|
"type": "fs",
|
|
"root": "/tmp",
|
|
"retry": {"max_attempts": 10},
|
|
}
|
|
]
|
|
}
|
|
)
|
|
src = cfg.sources[0]
|
|
assert src.retry is not None
|
|
assert src.retry.max_attempts == 10
|
|
# other retry fields fall back to the RetryPolicyConfig defaults
|
|
assert src.retry.base_delay_s == RetryPolicyConfig().base_delay_s
|
|
|
|
|
|
def test_yaml_round_trip():
|
|
yaml_text = """
|
|
ingester:
|
|
sources:
|
|
- type: fs
|
|
root: /data/docs
|
|
ignore_patterns: ["**/.git/**"]
|
|
delete_orphans: true
|
|
- type: s3
|
|
uri: s3://my-bucket/incoming/
|
|
poll_interval_s: 300
|
|
storage_options:
|
|
endpoint: http://seaweed:8333
|
|
- type: http
|
|
id: arxiv
|
|
urls: [https://arxiv.org/pdf/2301.12345.pdf]
|
|
headers:
|
|
Authorization: Bearer abc
|
|
poll_interval_s: 86400
|
|
workers:
|
|
worker_count: 8
|
|
max_concurrent: 4
|
|
api:
|
|
enabled: false
|
|
"""
|
|
data = yaml.safe_load(yaml_text)
|
|
app = AppConfig.model_validate(data)
|
|
assert len(app.ingester.sources) == 3
|
|
fs, s3, http = app.ingester.sources
|
|
assert isinstance(fs, FSSourceConfig)
|
|
assert isinstance(s3, S3SourceConfig)
|
|
assert isinstance(http, HTTPSourceConfig)
|
|
assert fs.ignore_patterns == ["**/.git/**"]
|
|
assert s3.storage_options["endpoint"] == "http://seaweed:8333"
|
|
assert http.headers["Authorization"] == "Bearer abc"
|
|
assert app.ingester.workers.worker_count == 8
|
|
assert app.ingester.api.enabled is False
|