File size: 4,235 Bytes
15d68eb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
"""
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.")