tests/test_generation_workspace.py

from __future__ import annotations

import json
import os
import shutil
import time
from pathlib import Path
from types import SimpleNamespace

os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")

import pytest
from PySide6.QtCore import QObject, QPoint, QPointF, QThreadPool, Qt, Signal
from PySide6.QtGui import QColor, QFontDatabase, QImage, QLinearGradient, QPainter, QPolygonF
from PySide6.QtWidgets import QApplication
from PySide6.QtTest import QTest

from adam.assets import AssetRegistry
from adam.config import ConfigManager
from adam.generations import GenerationRecord
from adam.registry import ToolRegistry
from adam.ui.asset_refresh import AssetRefresh
from adam.ui.generations import GenerationsPage
from adam.ui.theme import APP_STYLESHEET


def pump(app, condition=lambda: True):
    deadline = time.monotonic() + 5
    while time.monotonic() < deadline:
        app.processEvents()
        if condition():
            app.processEvents()
            return
        time.sleep(0.005)
    assert condition()


@pytest.fixture(scope="module")
def app():
    instance = QApplication.instance() or QApplication([])
    # Windows' offscreen Qt backend does not discover system fonts itself.
    for name in ("segoeui.ttf", "segoeuib.ttf", "seguisb.ttf", "seguisym.ttf"):
        font = Path("C:/Windows/Fonts") / name
        if font.exists():
            QFontDatabase.addApplicationFont(str(font))
    instance.setStyleSheet(APP_STYLESHEET)
    return instance


def landscape(path: Path, hue: int):
    """Synthetic review artwork; never reads the user's generation history."""
    image = QImage(640, 800, QImage.Format_RGB32)
    painter = QPainter(image)
    gradient = QLinearGradient(0, 0, 640, 800)
    gradient.setColorAt(0, QColor.fromHsv(hue, 170, 130))
    gradient.setColorAt(1, QColor.fromHsv((hue + 70) % 360, 130, 250))
    painter.fillRect(image.rect(), gradient)
    painter.setPen(Qt.NoPen)
    painter.setBrush(QColor("#ffe3a1"))
    painter.drawEllipse(390, 140, 120, 120)
    for offset, color in ((0, "#2e5274"), (130, "#153248"), (260, "#091b2d")):
        painter.setBrush(QColor(color))
        painter.drawPolygon(QPolygonF([QPointF(0, 520 + offset), QPointF(180, 290 + offset), QPointF(380, 580 + offset), QPointF(520, 410 + offset), QPointF(640, 570 + offset), QPointF(640, 800), QPointF(0, 800)]))
    painter.end()
    image.save(str(path))


@pytest.fixture
def workspace(app, tmp_path, monkeypatch):
    class Jobs(QObject):
        job_updated = Signal(object)

        def __init__(self):
            super().__init__()
            self.plans = []

        def submit(self, plan):
            self.plans.append(plan)
            return SimpleNamespace(id="layout-review")

    config = ConfigManager(tmp_path)
    shutil.copy2(Path(__file__).resolve().parents[1] / "config/tools.json", tmp_path / "config/tools.json")
    assets = AssetRegistry(tmp_path)
    base = tmp_path / "LandscapeXL.safetensors"
    base.touch()
    assets.register(kind="base_model", name="LandscapeXL", path=str(base))
    records = []
    for index, (trainer, provider, name) in enumerate((("lora", "lora_generator", "Mountain Light"), ("ddpm", "ddpm_generator", "Forest Study"), ("flow", "flow_generator", "Evening Coast"))):
        model = tmp_path / ("a_very_long_model_directory_name_" * 3) / (name + (".safetensors" if trainer == "lora" else ""))
        model.parent.mkdir(parents=True, exist_ok=True)
        if trainer == "lora":
            model.touch()
        else:
            model.mkdir()
            if trainer == "ddpm":
                (model / "model_index.json").write_text("{}")
            else:
                (model / "flow_model_info.json").write_text("{}")
                (model / "unet").mkdir()
                (model / "unet/config.json").write_text("{}")
        assets.register(kind="model", name=name, trainer=trainer, path=str(model))
        paths = []
        for number in range(6):
            image = tmp_path / f"image-{index}-{number}.png"
            landscape(image, 200 + index * 25 + number * 4)
            paths.append(str(image))
        metadata = tmp_path / f"generation-{index}.json"
        payload = dict(provider_id=provider, provider_name={"lora": "LoRA Generator", "ddpm": "DDPM Generator", "flow": "Flow Matching Generator"}[trainer],
                       model_name=name, model_path=str(model), images=paths, prompt="A quiet mountain valley at dusk, soft light and distant peaks",
                       negative_prompt="blur, watermark", base_model_path=str(base), seed=3210 + index * 10, steps=27,
                       sampler="DPM++ 2M" if trainer == "lora" else "DDIM" if trainer == "ddpm" else "Euler",
                       aspect_ratio="3:4 (Portrait)" if trainer == "lora" else "1:1 (Square)", width=864, height=1152,
                       cfg_scale=5.0, lora_strength=0.8, denoise_strength=0.5, created_at=f"2026-09-{26-index}T14:30:00+00:00")
        metadata.write_text(json.dumps(payload), encoding="utf-8")
        records.append(GenerationRecord.from_metadata(metadata))
    monkeypatch.setattr(AssetRefresh, "request", lambda *args, **kwargs: None)
    monkeypatch.setattr("adam.ui.generations.load_generation_history", lambda root: records)
    jobs = Jobs()
    page = GenerationsPage(tmp_path, ToolRegistry(tmp_path), jobs, assets, config)
    page.resize(1500, 900)
    page.show()
    pump(app, lambda: not page._history_reader._running and page.gallery.count() == 6)
    page.provider.setCurrentIndex(page.provider.findData("lora_generator"))
    page.base_model.setCurrentIndex(page.base_model.findData(str(base)))
    page.prompt.setPlainText(records[0].prompt)
    page.negative_prompt.setPlainText("blur, watermark")
    page.aspect.setCurrentText("3:4 (Portrait)")
    page.steps.setValue(27)
    page.cfg_scale.setValue(5.0)
    pump(app)
    yield page, jobs, records
    page.close()
    QThreadPool.globalInstance().waitForDone(3000)
    page.deleteLater()
    app.processEvents()


@pytest.mark.parametrize("width,height", [(1500, 900), (1200, 800)])
def test_layout_keeps_controls_inside_pane_and_preview_does_not_expand_it(app, workspace, width, height):
    page, _, _ = workspace
    page.resize(width, height)
    pump(app)
    assert page.size().width() == width
    assert page.size().height() == height
    create = page.workspace_splitter.widget(0)
    assert create.width() < width * 0.48
    assert page.workspace_splitter.widget(1).width() > create.width()
    assert create.horizontalScrollBar().maximum() == 0
    for control in (page.provider, page.model, page.base_model, page.width, page.height, page.images, page.preset, page.seed, page.generate_button):
        assert control.isVisible()
        right = control.mapTo(create.viewport(), QPoint(control.width(), 0)).x()
        assert right <= create.viewport().width(), control.objectName()
    assert page.advanced_toggle.isChecked()
    assert page.steps.isVisible() and page.cfg_scale.isVisible()
    if os.environ.get("ADAM_RENDER_GENERATION_REVIEW"):
        output = Path(__file__).resolve().parents[1] / "artifacts/generation-workspace-review"
        output.mkdir(parents=True, exist_ok=True)
        page.grab().save(str(output / f"generations-{width}x{height}.png"))


def test_dimensions_and_hidden_advanced_values_reach_generation_plan(app, workspace):
    page, jobs, _ = workspace
    page.width.setValue(888)
    page.height.setValue(1168)
    page.steps.setValue(39)
    page.cfg_scale.setValue(6.5)
    page.lora_strength.setValue(0.75)
    page.advanced_toggle.setChecked(False)
    page._generate()
    arguments = jobs.plans[-1].steps[0].arguments
    assert arguments["width"] == 888 and arguments["height"] == 1168
    assert arguments["steps"] == 39 and arguments["cfg_scale"] == 6.5
    assert arguments["lora_strength"] == 0.75
    saved = dict(page.config.get("generation_settings"))
    page.width.setValue(1024)
    page.config.update({"generation_settings": saved})
    page._restore_generation_settings()
    assert page.width.value() == 888 and page.height.value() == 1168


def test_history_tabs_navigation_favorites_and_comparison(app, workspace):
    page, _, records = workspace
    assert page.model_folders.count() == 3
    assert page._selection()[0] == records[0]
    QTest.mouseClick(page.model_folders.viewport(), Qt.LeftButton,
                     pos=page.model_folders.visualItemRect(page.model_folders.item(0)).center())
    pump(app)
    assert page.model_folders.count() == 1
    page.model_folders.setCurrentRow(0)
    QTest.keyClick(page.model_folders, Qt.Key_Return)
    pump(app, lambda: page.gallery.count() == 6)
    assert not page.model_folders.isVisible()
    page.history_tabs.setCurrentIndex(2)
    pump(app, lambda: page.gallery.count() == 18)
    page.gallery.setCurrentRow(7)
    selected_path = page._selection()[1]
    page._rate_selection("favorite")
    pump(app, lambda: not page._history_reader._running)
    page.history_tabs.setCurrentIndex(3)
    pump(app, lambda: page.gallery.count() == 1)
    assert page._selection()[1] == selected_path
    page.history_tabs.setCurrentIndex(4)
    pump(app, lambda: page.gallery.count() == 18)
    page.gallery.setCurrentRow(0)
    page._pin_comparison()
    pinned = page.comparison_caption.text()
    page.gallery.setCurrentRow(1)
    assert page.comparison_panel.isVisible()
    assert page.comparison_caption.text() == pinned
    page._clear_displayed_history()
    pump(app)
    assert page.gallery.count() == 0


@pytest.mark.parametrize("provider", ["ddpm_generator", "flow_generator"])
def test_provider_specific_dimensions_and_smart_settings_survive_collapsing(app, workspace, provider):
    page, jobs, _ = workspace
    page.provider.setCurrentIndex(page.provider.findData(provider))
    ddpm = provider == "ddpm_generator"
    custom = page.ddpm_custom_size if ddpm else page.flow_custom_size
    width = page.ddpm_width if ddpm else page.flow_width
    height = page.ddpm_height if ddpm else page.flow_height
    custom.setChecked(True)
    width.setValue(384)
    height.setValue(512)
    page.smart_enabled.setChecked(True)
    page.advanced_toggle.setChecked(False)
    page._generate()
    arguments = jobs.plans[-1].steps[0].arguments
    assert arguments["width"] == 384 and arguments["height"] == 512
    if "smart_generation" in page._current_tool().capabilities:
        assert arguments["smart_generation"] is True
    assert not page.lora_dimensions.isVisible()


def test_reusing_lora_image_restores_dimensions_and_strengths(workspace):
    page, _, _ = workspace
    page.width.setValue(1024)
    page.height.setValue(1024)
    page.cfg_scale.setValue(10.0)
    page.lora_strength.setValue(1.5)
    page.negative_prompt.clear()
    page._reuse_settings()
    assert (page.width.value(), page.height.value()) == (864, 1152)
    assert page.cfg_scale.value() == 5.0
    assert page.lora_strength.value() == 0.8
    assert page.negative_prompt.toPlainText() == "blur, watermark"