haiku.rag/tests/ingester/test_config.py

from pathlib import Path

import pytest
import yaml
from pydantic import ValidationError

from haiku.rag.config import (
    APIConfig,
    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
    assert cfg.api.root_path == ""


@pytest.mark.parametrize(
    "raw, expected",
    [
        ("", ""),
        ("/", ""),
        ("ingester", "/ingester"),
        ("/ingester", "/ingester"),
        ("/ingester/", "/ingester"),
        ("  /ingester/  ", "/ingester"),
    ],
)
def test_api_root_path_normalized(raw, expected):
    cfg = APIConfig(root_path=raw)

    assert cfg.root_path == expected


def test_api_root_path_normalized_on_assignment():
    """validate_assignment ensures CLI overrides (--root-path) get the same
    normalization as values parsed from the config file."""
    cfg = APIConfig()

    cfg.root_path = "ingester/"

    assert cfg.root_path == "/ingester"


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
    assert cfg.sources[0].metadata_provider is None


def test_source_carries_metadata_provider_name():
    cfg = IngesterConfig.model_validate(
        {
            "sources": [
                {"type": "fs", "root": "/tmp/docs", "metadata_provider": "example"}
            ]
        }
    )
    assert cfg.sources[0].metadata_provider == "example"


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
  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