AareDAQ/tests/unit/gui/test_models.py

import pytest
from aarecommon.models.models import DewarAddress, SampleShortInfo
from PySide6.QtCore import Qt

from aare.gui.models.sample_queue_model import SampleQueueSpreadsheet
from aare.gui.models.user_sample_model import UserSampleSpreadsheet


@pytest.fixture
def sample_list():
    return [
        SampleShortInfo(
            db_id=1,
            puck_name="P1",
            dewar_name="D1",
            sample_name="S1",
            run_number=1,
            user="U1",
            pin=1,
            location=DewarAddress(segment="A", pos=1),
        ),
        SampleShortInfo(
            db_id=2,
            puck_name="P2",
            dewar_name="D2",
            sample_name="S2",
            run_number=2,
            user="U2",
            pin=2,
            location=DewarAddress(segment="A", pos=2),
        ),
        SampleShortInfo(
            db_id=3,
            puck_name="P1",
            dewar_name="D1",
            sample_name="S3",
            run_number=3,
            user="U1",
            pin=3,
            location=DewarAddress(segment="A", pos=1),
        ),
    ]


def test_user_sample_model_init(sample_list):
    model = UserSampleSpreadsheet(samples=sample_list)
    model.set_show_all_pgroups(True)  # Ensure all pgroups are shown for testing
    assert model.rowCount() == 3
    # Check if filtering works
    # "User" is column 5
    model.set_filter("User", "U2")
    assert model.rowCount() == 1
    model.clear_filter()
    assert model.rowCount() == 3


def test_user_sample_model_column_filter(sample_list):
    model = UserSampleSpreadsheet(samples=sample_list)
    model.set_show_all_pgroups(True)
    # Column 6 is User (column 0 is the frozen #+status cell)
    model.set_column_filter(6, "U1")
    assert model.rowCount() == 2
    model.clear_all_column_filters()
    assert model.rowCount() == 3


def test_user_sample_model_unique_values(sample_list):
    model = UserSampleSpreadsheet(samples=sample_list)
    model.set_show_all_pgroups(True)
    # Column 6 is User (column 0 is the frozen #+status cell)
    users = model.unique_values_for_column(6)
    assert "U1" in users
    assert "U2" in users
    assert len(users) == 2


def test_sample_queue_model_init(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    assert model.rowCount() == 2
    assert model.columnCount() == 3
    assert model.data(model.index(0, 0), Qt.ItemDataRole.DisplayRole) == "D1"
    assert model.data(model.index(1, 2), Qt.ItemDataRole.DisplayRole) == "S2"


def test_sample_queue_model_update(sample_list):
    model = SampleQueueSpreadsheet()
    assert model.rowCount() == 0
    model.updateData(sample_list[:2])
    assert model.rowCount() == 2


def test_sample_queue_model_remove(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    model.remove_sample(1)
    assert model.rowCount() == 1
    assert model.samples[0].db_id == 2


def test_sample_queue_model_clear(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    model.clearSamples()
    assert model.rowCount() == 0


def test_sample_queue_model_set_running(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    # Background color role for first row
    color_not_running = model.data(model.index(0, 0), Qt.ItemDataRole.BackgroundRole)
    model.set_running(True)
    color_running = model.data(model.index(0, 0), Qt.ItemDataRole.BackgroundRole)
    assert color_not_running != color_running


def test_sample_queue_model_header(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    assert model.headerData(0, Qt.Orientation.Horizontal, Qt.ItemDataRole.DisplayRole) == "Dewar"
    assert model.headerData(1, Qt.Orientation.Vertical, Qt.ItemDataRole.DisplayRole) == "2"


def test_sample_queue_model_flags(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    flags = model.flags(model.index(0, 0))
    assert flags & Qt.ItemFlag.ItemIsDropEnabled


# --- Status logic of the combined dewar/queue view ---------------------------
# The dewar table doubles as the queue view: the frozen "#" column carries a
# status fill (mounted > queued > flagged > measured) and the chip row filters
# by status. This is the logic a local contact trusts at a glance, so it gets
# its own tests.


def _status(model, row):
    brush = model.data(model.index(row, 0), Qt.ItemDataRole.BackgroundRole)
    return None if brush is None else brush.color().name().lower()


def _row_of(model, db_id):
    return next(r for r in range(model.rowCount()) if model.get_id(r).db_id == db_id)


@pytest.fixture
def status_model(sample_list):
    from aarecommon.models.models import DewarAddress, SampleShortInfo

    # A measured sample: any rotation data counts (exactly 1 MUST count;
    # unmeasured is rotation_count 0, like the fixture's samples 1-3).
    sample_list.append(
        SampleShortInfo(
            db_id=4,
            puck_name="P3",
            dewar_name="D3",
            sample_name="S4",
            run_number=4,
            user="U1",
            pin=4,
            rotation_count=1,
            location=DewarAddress(segment="B", pos=1),
        )
    )
    model = UserSampleSpreadsheet(samples=sample_list)
    model.set_show_all_pgroups(True)
    return model


def test_status_color_priority(status_model):
    from aare.gui.styles import (
        SAMPLE_ROW_QUEUED_BG,
        SAMPLE_STATUS_FLAGGED_BG,
        SAMPLE_STATUS_MEASURED_BG,
        SAMPLE_STATUS_QUEUED_BG,
    )

    model = status_model
    assert _status(model, _row_of(model, 1)) is None

    model.set_queued_ids({1})
    model.set_flagged(1, True)
    # Queued beats flagged in the All view.
    assert _status(model, _row_of(model, 1)) == SAMPLE_STATUS_QUEUED_BG.lower()
    model.set_queued_ids(set())
    assert _status(model, _row_of(model, 1)) == SAMPLE_STATUS_FLAGGED_BG.lower()

    # Measured is automatic: rotation_count 1 counts, the 0s of 1-3 don't.
    assert _status(model, _row_of(model, 4)) == SAMPLE_STATUS_MEASURED_BG.lower()
    assert _status(model, _row_of(model, 2)) is None

    # Mounted always wins.
    model.updateCurrentSample(current_puck="P1", current_sample=1)
    assert _status(model, _row_of(model, 1)) == SAMPLE_ROW_QUEUED_BG.lower()


def test_status_filter_selects_rows(status_model):
    model = status_model
    model.set_queued_ids({1, 2})
    model.set_flagged(3, True)

    model.set_status_filter("queued")
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {1, 2}
    model.set_status_filter("flagged")
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {3}
    model.set_status_filter("measured")
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {4}
    # Unmeasured is the complement: everything still to be done.
    model.set_status_filter("unmeasured")
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {1, 2, 3}
    model.set_status_filter(None)
    assert model.rowCount() == 4


def test_queued_view_shows_queue_order(status_model):
    model = status_model
    # Queue order deliberately different from location/db order.
    model.set_queued_ids([3, 1, 2])
    model.set_status_filter("queued")
    assert [model.get_id(r).db_id for r in range(model.rowCount())] == [3, 1, 2]
    # Header sorts must not reorder the queue view — row 1 runs next.
    model.sort(1, Qt.SortOrder.AscendingOrder)
    assert [model.get_id(r).db_id for r in range(model.rowCount())] == [3, 1, 2]
    # Other views keep the normal header sort.
    model.set_status_filter(None)
    assert [model.get_id(r).db_id for r in range(model.rowCount())] != [3, 1, 2]


def test_status_tints_are_context_dependent(status_model):
    from aare.gui.styles import SAMPLE_STATUS_FLAGGED_BG, SAMPLE_STATUS_QUEUED_BG

    model = status_model
    model.set_queued_ids({1, 2})
    model.set_flagged(1, True)

    # Queued view: own tint suppressed, only the also-flagged mark shows.
    model.set_status_filter("queued")
    assert _status(model, _row_of(model, 1)) == SAMPLE_STATUS_FLAGGED_BG.lower()
    assert _status(model, _row_of(model, 2)) is None

    # Flagged view: a re-queued sample wears the queued mark.
    model.set_status_filter("flagged")
    assert _status(model, _row_of(model, 1)) == SAMPLE_STATUS_QUEUED_BG.lower()


def test_status_sets_refilter_while_chip_active(status_model):
    model = status_model
    model.set_status_filter("queued")
    assert model.rowCount() == 0
    model.set_queued_ids({2})
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {2}


def test_status_column_is_display_only(status_model):
    model = status_model
    assert model.data(model.index(0, 0), Qt.ItemDataRole.DisplayRole) == 1
    before = [model.get_id(r).db_id for r in range(model.rowCount())]
    model.sort(0, Qt.SortOrder.DescendingOrder)  # no-op on the "#" column
    assert [model.get_id(r).db_id for r in range(model.rowCount())] == before


def test_mime_data_round_trips_for_chip_drops(status_model):
    from aarecommon.models.models import SampleShortInfoList

    model = status_model
    payload = model.mimeData([model.index(0, 1), model.index(1, 1)])
    samples = SampleShortInfoList.model_validate_json(payload.text())
    assert len(samples.s) == 2
    assert samples.s[0].db_id == model.get_id(0).db_id


def test_mounted_row_is_blue_across_all_columns(status_model):
    # Whole-row blue (not just "#") so the mounted sample is findable at a
    # glance; the blue keeps the theme's own ink (no dark override).
    model = status_model
    model.updateCurrentSample(current_puck="P1", current_sample=1)
    row = _row_of(model, 1)
    for col in range(model.columnCount()):
        brush = model.data(model.index(row, col), Qt.ItemDataRole.BackgroundRole)
        assert brush.color().name().lower() == "#729fcf"
        assert model.data(model.index(row, col), Qt.ItemDataRole.ForegroundRole) is None
    # Other rows: still "#"-only tint, other columns untouched.
    other = _row_of(model, 2)
    assert model.data(model.index(other, 1), Qt.ItemDataRole.BackgroundRole) is None


def test_spreadsheet_param_columns(status_model):
    from aarecommon.models.models import DataCollectionParameters

    model = status_model
    hdr = model.header
    osc, exp, tot, trans = (
        hdr.index("Oscillation (°)"),
        hdr.index("Exposure (s)"),
        hdr.index("Total range (°)"),
        hdr.index("Transmission"),
    )
    model.get_id(_row_of(model, 1)).aaredb_params = DataCollectionParameters(
        oscillation=0.1, exposure=0.01, totalangle=360, transmission=0.125
    )
    row = _row_of(model, 1)
    d = lambda c: model.data(model.index(row, c), Qt.ItemDataRole.DisplayRole)
    # Transmission shown as stored, no percent conversion.
    assert (d(osc), d(exp), d(tot), d(trans)) == (0.1, 0.01, 360, 0.125)
    # No params -> blank, and sorting a mostly-None column must not raise;
    # the real value sorts first.
    assert model.data(model.index(_row_of(model, 2), osc), Qt.ItemDataRole.DisplayRole) is None
    model.sort(osc, Qt.SortOrder.AscendingOrder)
    assert model.get_id(0).db_id == 1
    assert model.headerData(hdr.index("Comment"), Qt.Orientation.Horizontal, 0) == "Comment"