tests/test_agents.py

"""
Smoke tests for Indic Heritage Studio v2.

Run with:
    pytest tests/test_agents.py -v

These tests don't require a GPU — they verify the agent layer, config,
and utility functions. GPU pipeline tests live in test_pipelines.py.
"""
from __future__ import annotations

import os
import sys
from pathlib import Path

# Ensure project root is on path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))


def test_settings_loads():
    from config.settings import settings
    assert settings.t2i_model_id.startswith("lykon/") or "sdxl" in settings.t2i_model_id.lower()
    assert settings.default_image_size == 1024
    assert settings.svd_num_frames == 25
    assert settings.torch_dtype_str == "float16"


def test_styles_complete():
    from config.styles import HERITAGE_STYLES, list_styles, get_style
    assert len(HERITAGE_STYLES) == 5
    for s in list_styles():
        assert s.id in {"madhubani", "warli", "pattachitra", "mughal", "tanjore"}
        assert len(s.prompt_tags) >= 5
        assert len(s.palette) >= 4
        assert 0 < s.lora_scale <= 1.0
        assert 1 <= s.svd_motion_bucket <= 255


def test_style_advisor_fallback():
    """StyleAdvisor should return a valid recommendation even without API key."""
    from agents.style_advisor import StyleAdvisor
    adv = StyleAdvisor()
    # Force fallback path
    adv.client._api_key = ""
    result = adv.recommend("a courtly scene with a king and ministers")
    assert "style" in result
    assert result["style"] in {"madhubani", "warli", "pattachitra", "mughal", "tanjore"}
    assert result["source"] == "heuristic_fallback"


def test_prompt_engineer_fallback():
    from agents.prompt_engineer import PromptEngineer
    from config.styles import get_style
    eng = PromptEngineer()
    eng.client._api_key = ""
    style = get_style("madhubani")
    result = eng.enrich("a woman reading", style)
    assert result.ok
    assert "madhubani" in result.content.lower()


def test_prompt_engineer_negative():
    from agents.prompt_engineer import PromptEngineer
    from config.styles import get_style
    neg = PromptEngineer.build_negative(get_style("warli"))
    assert "photorealistic" in neg
    assert "low quality" in neg


def test_critic_heuristic():
    from agents.critic import Critic
    from config.styles import get_style
    from PIL import Image
    import numpy as np
    # Random RGB image
    arr = np.random.randint(0, 256, (512, 512, 3), dtype=np.uint8)
    img = Image.fromarray(arr)
    critic = Critic()
    critic.client._api_key = ""
    result = critic.evaluate(img, get_style("madhubani"), "test prompt")
    assert 1 <= result.style_fidelity <= 10
    assert 1 <= result.composition <= 10
    assert 1 <= result.technical_quality <= 10
    assert 0 <= result.overall <= 10
    assert result.source == "heuristic"


def test_gpu_utils_shard():
    from utils.gpu_utils import shard_workload
    items = list(range(10))
    shards = shard_workload(items, 3)
    assert len(shards) == 3
    assert sum(len(s) for s in shards) == 10
    # First shards should be larger by 1 if remainder
    assert len(shards[0]) == 4
    assert len(shards[1]) == 3
    assert len(shards[2]) == 3


def test_batch_job_builder():
    from core.batch_processor import BatchProcessor
    jobs = BatchProcessor.build_t2i_jobs(
        prompt="test", styles=["madhubani", "warli"],
        seeds=[1, 2], output_dir=Path("/tmp/test_batch"),
    )
    assert len(jobs) == 4  # 2 styles × 2 seeds
    assert all(j.mode == "t2i" for j in jobs)
    assert all(j.prompt == "test" for j in jobs)


def test_image_utils_resize():
    from utils.image_utils import resize_to_sdxl
    from PIL import Image
    img = Image.new("RGB", (800, 600), "red")
    out = resize_to_sdxl(img, target=1024)
    assert out.size == (1024, 768)  # 4:3 from 800x600


if __name__ == "__main__":
    # Run as a script if pytest isn't available
    test_settings_loads()
    test_styles_complete()
    test_style_advisor_fallback()
    test_prompt_engineer_fallback()
    test_prompt_engineer_negative()
    test_critic_heuristic()
    test_gpu_utils_shard()
    test_batch_job_builder()
    test_image_utils_resize()
    print("All tests passed.")