tests/test_ui_loading.py

from __future__ import annotations

import os
import shutil
import threading
import time
from copy import deepcopy
from dataclasses import asdict
from pathlib import Path

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

import pytest
from PySide6.QtCore import QCoreApplication, QThreadPool, QTimer, QObject, Signal
from PySide6.QtWidgets import QApplication, QStackedWidget, QWidget

from adam.assets import AssetRegistry
from adam.config import ConfigManager
from adam.experiment_tracker import ExperimentStore
from adam.ui.asset_refresh import AssetRefresh
from adam.ui.background import BackgroundRead
from adam.ui.model_intelligence import ModelIntelligencePage
from adam.ui.main_window import (
    expand_subject_reference,
    is_explicit_action_request,
    is_pending_job_confirmation_request,
    subject_idea_from_chat,
)


@pytest.fixture(scope="module")
def app():
    instance = QCoreApplication.instance()
    if instance is not None and not isinstance(instance, QApplication):
        pytest.skip("Run UI loading tests separately from tests that create QCoreApplication.")
    return instance or QApplication([])


def pump(app, condition, timeout=5):
    deadline = time.monotonic() + timeout
    while not condition() and time.monotonic() < deadline:
        app.processEvents()
        time.sleep(0.005)
    assert condition()


def test_background_read_keeps_navigation_responsive_and_delivers_on_ui_thread(app):
    stack = QStackedWidget()
    stack.addWidget(QWidget())
    stack.addWidget(QWidget())
    reader = BackgroundRead(stack)
    gate = threading.Event()
    started = threading.Event()
    results = []
    reader.ready.connect(lambda value: results.append((value, threading.get_ident())))

    def slow_read():
        started.set()
        assert gate.wait(3)
        return "loaded"

    try:
        reader.request(slow_read)
        pump(app, started.is_set)
        QTimer.singleShot(0, lambda: stack.setCurrentIndex(1))
        pump(app, lambda: stack.currentIndex() == 1)
        assert results == []
    finally:
        gate.set()
    pump(app, lambda: bool(results))
    assert results == [("loaded", threading.get_ident())]


def test_chat_action_requests_are_routed_to_the_validated_planner() -> None:
    assert is_explicit_action_request("Can you train a DDPM on cozy cafe images?")
    assert is_explicit_action_request("Yes, start this job now")
    assert not is_explicit_action_request("What subject should we train today?")
    assert not is_explicit_action_request("Explain how DDPM training works")
    assert is_pending_job_confirmation_request("Yes, start this job now")
    assert not is_pending_job_confirmation_request("Can you train a DDPM on cozy cafe images?")


def test_subject_reference_is_expanded_from_the_prior_chat_idea() -> None:
    reply = """🎯 Subject Idea:
> “Astronauts Exploring a Floating City in Jupiter's Atmosphere”
"""
    assert subject_idea_from_chat(reply) == "Astronauts Exploring a Floating City in Jupiter's Atmosphere"
    request, subject = expand_subject_reference("Train a DDPM with that subject you mentioned", reply)
    assert subject == "Astronauts Exploring a Floating City in Jupiter's Atmosphere"
    assert request == "Train a DDPM with Astronauts Exploring a Floating City in Jupiter's Atmosphere"


def test_refresh_coalesces_requests_and_does_not_publish_stale_data(app):
    owner = QWidget()
    reader = BackgroundRead(owner)
    gate = threading.Event()
    results = []
    reader.ready.connect(results.append)
    try:
        reader.request(lambda: gate.wait(3) and "old")
        reader.request(lambda: "superseded", force=True)
        reader.request(lambda: "latest", force=True)
    finally:
        gate.set()
    pump(app, lambda: bool(results))
    assert results == ["latest"]
    reader.request(lambda: "unnecessary repeat")
    app.processEvents()
    assert not reader._running
    assert results == ["latest"]


def test_failed_read_can_be_retried(app):
    owner = QWidget()
    reader = BackgroundRead(owner)
    errors, results = [], []
    reader.failed.connect(errors.append)
    reader.ready.connect(results.append)

    def fail():
        raise OSError("disconnected")

    reader.request(fail)
    pump(app, lambda: bool(errors))
    reader.request(lambda: "recovered")
    pump(app, lambda: bool(results))
    assert errors == ["disconnected"]
    assert results == ["recovered"]


def test_asset_merge_preserves_changes_made_during_discovery(app, tmp_path):
    assets = AssetRegistry(tmp_path)
    original = assets.register(kind="model", name="Original", path=str(tmp_path / "original"))
    removed = assets.register(kind="model", name="Remove me", path=str(tmp_path / "removed"))
    baseline = {item.id: asdict(item) for item in assets.assets}
    discovered = deepcopy(assets.assets)
    original.name = "User edited name"
    assets.assets.remove(removed)
    new = assets.register(kind="model", name="Just trained", path=str(tmp_path / "new"))
    service = AssetRefresh(assets, ConfigManager(tmp_path))
    service._merge((baseline, discovered))
    assert {item.name for item in assets.assets} == {"User edited name", "Just trained"}
    assert assets.assets[0] is original
    assert new in assets.assets


def test_background_discovery_does_not_write_shared_state(app, tmp_path, monkeypatch):
    assets = AssetRegistry(tmp_path)
    config = ConfigManager(tmp_path)
    writes = []
    monkeypatch.setattr(AssetRegistry, "save", lambda self: writes.append(threading.get_ident()))
    monkeypatch.setattr("adam.dataset_registry.DatasetRegistry.save", lambda self: writes.append(threading.get_ident()))
    assets.discover(config, persist=False)
    assert writes == []


def test_intelligence_load_does_not_block_other_tabs(app, tmp_path, monkeypatch):
    gate = threading.Event()
    started = threading.Event()

    def history(*args, **kwargs):
        started.set()
        assert gate.wait(3)
        return []

    monkeypatch.setattr("adam.ui.model_intelligence.load_generation_history", history)
    page = ModelIntelligencePage(tmp_path, ExperimentStore(tmp_path))
    stack = QStackedWidget()
    stack.addWidget(page)
    stack.addWidget(QWidget())
    try:
        page.refresh()
        pump(app, started.is_set)
        QTimer.singleShot(0, lambda: stack.setCurrentIndex(1))
        pump(app, lambda: stack.currentIndex() == 1)
        assert page._reader._running
    finally:
        gate.set()
    pump(app, lambda: not page._reader._running)
    assert page.metric_models.text() == "0"
    QThreadPool.globalInstance().waitForDone(3000)


def test_generation_refresh_keeps_edits_and_folder_navigation_reuses_history(app, tmp_path, monkeypatch):
    from adam.registry import ToolRegistry
    from adam.ui.generations import GenerationsPage

    class Jobs(QObject):
        job_updated = Signal(object)

    config = ConfigManager(tmp_path)
    shutil.copy2(Path(__file__).resolve().parents[1] / "config/tools.json", tmp_path / "config/tools.json")
    assets = AssetRegistry(tmp_path)
    monkeypatch.setattr(AssetRefresh, "request", lambda *args, **kwargs: None)
    calls = []

    def history(*args, **kwargs):
        calls.append(1)
        return []

    monkeypatch.setattr("adam.ui.generations.load_generation_history", history)
    jobs = Jobs()
    page = GenerationsPage(tmp_path, ToolRegistry(tmp_path), jobs, assets, config)
    pump(app, lambda: not page._history_reader._running)
    page.provider.setCurrentIndex(page.provider.findData("ddpm_generator"))
    page.steps.setValue(37)
    page.seed.setText("123456")  # In-progress edit has not emitted editingFinished.
    page.prompt.setPlainText("Keep my current prompt")
    model = tmp_path / "model"
    model.mkdir()
    (model / "model_index.json").write_text("{}")
    assets.register(kind="model", name="New model", path=str(model), trainer="ddpm")
    page.refresh()
    assert page.steps.value() == 37
    assert page.seed.text() == "123456"
    assert page.prompt.toPlainText() == "Keep my current prompt"
    assert page.model.currentData() == str(model)
    assert page.generate_button.isEnabled()
    page._show_provider_folders()
    page._go_back_in_history()
    app.processEvents()
    assert calls == [1]