tests/test_workbench.py

from __future__ import annotations

import json

from kneiff.app.workbench import (
    EXAMPLE_TAGS,
    EXAMPLE_VALIDATION_PROMPT_CONFIG,
    PROFILE_CHROMA_LABEL,
    PROFILE_NLG_LABEL,
    build_demo,
    generate_validation_prompts,
    render_tags,
)

NEGATIVE_CONTROL_KEYS = (
    "negative_control_wolf",
    "negative_control_human_residential_street",
    "negative_control_human_office_worker",
)


def test_render_tags_defaults_to_existing_hybrid_caption() -> None:
    caption, diagnostics = render_tags(EXAMPLE_TAGS)

    assert caption.startswith("Character_Token, species_token, solo, front_view")
    assert (
        "\nThis realistic image shows a standing character in a full-body front-view."
        in caption
    )
    assert "realistic" in caption
    assert "Validation: ok" in diagnostics


def test_render_tags_renders_nlg_zimage_profile() -> None:
    caption, diagnostics = render_tags(EXAMPLE_TAGS, PROFILE_NLG_LABEL)

    assert caption.startswith(
        "This realistic image shows a standing character in a full-body front-view."
    )
    assert "The character is looking to the side" in caption
    assert "character:" not in caption
    assert "species:" not in caption
    assert "Validation: ok" in diagnostics


def test_render_tags_renders_chroma_profile() -> None:
    caption, diagnostics = render_tags(EXAMPLE_TAGS, PROFILE_CHROMA_LABEL)

    assert caption.startswith("identity_reference. full_body_reference.")
    assert "character:Character_Token." in caption
    assert "species:species_token." in caption
    assert "Validation: ok" in diagnostics


def test_render_tags_reports_unknown_profile() -> None:
    caption, diagnostics = render_tags(EXAMPLE_TAGS, "Nope")

    assert caption == ""
    assert "Validation: Unknown profile: Nope" in diagnostics


def test_render_tags_preserves_unknown_inputs() -> None:
    caption, diagnostics = render_tags("solo, workbench_unknown")

    assert "workbench_unknown" in caption
    assert "Unknown tags preserved: workbench_unknown" in diagnostics


def test_generate_validation_prompts_renders_controls_and_subsets() -> None:
    prompt_json, diagnostics = generate_validation_prompts(
        EXAMPLE_VALIDATION_PROMPT_CONFIG
    )

    prompts = json.loads(prompt_json)

    assert list(prompts) == [
        *NEGATIVE_CONTROL_KEYS,
        "activation_control",
        "subset_expression",
        "subset_identity",
    ]
    assert "male wolf character" in prompts["negative_control_wolf"].lower()
    assert prompts["negative_control_human_residential_street"].startswith(
        "A young human man walks alone"
    )
    assert prompts["negative_control_human_office_worker"].startswith(
        "A tired human office worker sits"
    )
    assert prompts["activation_control"].startswith("Character_Token.")
    assert "smiling" in prompts["subset_expression"]
    assert "front-view portrait" in prompts["subset_expression"]
    assert "looking at the viewer and smiling" in prompts["subset_expression"]
    assert "Prompt count: 6" in diagnostics
    assert "Validation: ok" in diagnostics


def test_generate_validation_prompts_reports_invalid_yaml() -> None:
    prompt_json, diagnostics = generate_validation_prompts(
        "profile: nlg\nprofile: chroma\n"
    )

    assert prompt_json == ""
    assert "Validation: Duplicate YAML key 'profile'" in diagnostics


def test_generate_validation_prompts_reports_unknown_tags() -> None:
    prompt_json, diagnostics = generate_validation_prompts(
        EXAMPLE_VALIDATION_PROMPT_CONFIG.replace("smile", "not_a_real_tag")
    )

    assert prompt_json == ""
    assert "unknown tags: not_a_real_tag" in diagnostics


def test_generate_validation_prompts_allows_warn_policy_axis_conflicts() -> None:
    prompt_json, diagnostics = generate_validation_prompts(
        EXAMPLE_VALIDATION_PROMPT_CONFIG.replace(
            "front_view, full_body",
            "front_view, side_view, full_body",
        )
    )

    assert json.loads(prompt_json)
    assert "Validation: ok" in diagnostics


def test_build_demo_constructs_blocks() -> None:
    demo = build_demo()

    assert demo is not None