scheduler/backend/test_bundle.py

from __future__ import annotations

from dataclasses import dataclass
import os
from test_config import build_bundle_test_configs


@dataclass(frozen=True)
class Test:
    test_id: str
    device: str
    test_type: str
    rotation: str | None
    rx_tx: str | None
    has_coe_pair: bool
    coe_pairing: list[str]
    config: dict[str, dict[str, str | None]]
    throttled: bool
    estimated_minutes: int

# Bundle priority tiers for scheduling order (lower number = higher priority)
BUNDLE_PRIORITY_TOP = 0  # Failed tests requiring rerun
BUNDLE_PRIORITY_P2P_WITH_COE = 1  # P2P tests with COE pairs
BUNDLE_PRIORITY_P2P_ONLY = 2  # P2P tests without COE pairs (RX/TX bundled)
BUNDLE_PRIORITY_COE_ONLY = 3  # COE tests without P2P pairing (should be rare/unschedulable)
BUNDLE_PRIORITY_P3P = 4  # P3P tests

DUT = (os.getenv("DUT") or "DUT").strip()
REF = (os.getenv("REF") or "REF").strip()

@dataclass(frozen=True)
class TestBundle:
    index: int
    tests: list[str]
    total_minutes: int
    device: str
    config: tuple[str, ...]
    priority: int
    
def build_test_bundles(active_dut: dict[str, Test], active_ref: dict[str, Test], top_priority_tests: set[tuple[str, str]]) -> list[TestBundle]:
    """Build deterministic bundles for DP scheduling."""

    processed_dut: set[str] = set()
    processed_ref: set[str] = set()
    test_bundles: list[TestBundle] = []
    bundle_index = 0

    def dedupe_preserve_order(test_ids: list[str]) -> list[str]:
        seen: set[str] = set()
        deduped: list[str] = []
        for test_id in test_ids:
            if test_id in seen:
                continue
            seen.add(test_id)
            deduped.append(test_id)
        return deduped

    dut_top_priority: list[str] = []
    ref_top_priority: list[str] = []

    for test_id, device in top_priority_tests:
        if device == DUT and test_id in active_dut:
            dut_top_priority.append(test_id)
        elif device == REF and test_id in active_ref:
            ref_top_priority.append(test_id)

    if dut_top_priority:
        test_bundles.append(create_bundle(dut_top_priority, DUT, BUNDLE_PRIORITY_TOP, bundle_index, active_dut, active_ref))
        processed_dut.update(dut_top_priority)
        bundle_index += 1

    if ref_top_priority:
        test_bundles.append(create_bundle(ref_top_priority, REF, BUNDLE_PRIORITY_TOP, bundle_index, active_dut, active_ref))
        processed_ref.update(ref_top_priority)
        bundle_index += 1

    # Phase 1: DUT P2P bundles
    for dut_test_id, dut_test in active_dut.items():
        if dut_test_id in processed_dut or dut_test.test_type != "P2P":
            continue

        dut_bundled_tests = [dut_test_id]
        processed_dut.add(dut_test_id)

        active_coe_pairings = [
            test_id
            for test_id in dut_test.coe_pairing
            if test_id in active_dut and test_id not in processed_dut
        ]

        if active_coe_pairings:
            dut_bundled_tests.extend(active_coe_pairings)
            processed_dut.update(active_coe_pairings)
            priority = BUNDLE_PRIORITY_P2P_WITH_COE
        else:
            rx_tx_pair_id = get_rx_tx_pair_id(dut_test, DUT, active_dut, active_ref)
            if rx_tx_pair_id and rx_tx_pair_id not in processed_dut:
                dut_bundled_tests.append(rx_tx_pair_id)
                processed_dut.add(rx_tx_pair_id)
            priority = BUNDLE_PRIORITY_P2P_ONLY

        dut_bundled_tests = dedupe_preserve_order(dut_bundled_tests)
        test_bundles.append(
            create_bundle(dut_bundled_tests, DUT, priority, bundle_index, active_dut, active_ref)
        )

        ref_bundled_tests = [
            test_id
            for test_id in dut_bundled_tests
            if test_id in active_ref and test_id not in processed_ref
        ]
        if ref_bundled_tests:
            ref_bundled_tests = dedupe_preserve_order(ref_bundled_tests)
            processed_ref.update(ref_bundled_tests)
            test_bundles.append(
                create_bundle(ref_bundled_tests, REF, priority, bundle_index, active_dut, active_ref)
            )

        bundle_index += 1

    # Phase 2: DUT P3P bundles
    for dut_test_id, dut_test in active_dut.items():
        if dut_test_id in processed_dut or dut_test.test_type != "P3P":
            continue

        dut_bundled_tests = [dut_test_id]
        processed_dut.add(dut_test_id)

        th_ut_pair_id = get_th_ut_pair_id(dut_test, DUT, active_dut, active_ref)
        if th_ut_pair_id and th_ut_pair_id not in processed_dut:
            dut_bundled_tests.append(th_ut_pair_id)
            processed_dut.add(th_ut_pair_id)

        dut_bundled_tests = dedupe_preserve_order(dut_bundled_tests)
        test_bundles.append(
            create_bundle(dut_bundled_tests, DUT, BUNDLE_PRIORITY_P3P, bundle_index, active_dut, active_ref)
        )

        ref_bundled_tests = [
            test_id
            for test_id in dut_bundled_tests
            if test_id in active_ref and test_id not in processed_ref
        ]
        if ref_bundled_tests:
            ref_bundled_tests = dedupe_preserve_order(ref_bundled_tests)
            processed_ref.update(ref_bundled_tests)
            test_bundles.append(
                create_bundle(ref_bundled_tests, REF, BUNDLE_PRIORITY_P3P, bundle_index, active_dut, active_ref)
            )

        bundle_index += 1

    # Phase 3: DUT leftovers (including COE-only)
    for dut_test_id in active_dut:
        if dut_test_id in processed_dut:
            continue
        test_bundles.append(
            create_bundle([dut_test_id], DUT, BUNDLE_PRIORITY_COE_ONLY, bundle_index, active_dut, active_ref)
        )
        bundle_index += 1
        processed_dut.add(dut_test_id)

    # Phase 4: unmatched REF P2P bundles
    for ref_test_id, ref_test in active_ref.items():
        if ref_test_id in processed_ref or ref_test.test_type != "P2P":
            continue

        ref_bundled_tests = [ref_test_id]
        processed_ref.add(ref_test_id)

        active_coe_pairings = [
            test_id
            for test_id in ref_test.coe_pairing
            if test_id in active_ref and test_id not in processed_ref
        ]

        if active_coe_pairings:
            ref_bundled_tests.extend(active_coe_pairings)
            processed_ref.update(active_coe_pairings)
            priority = BUNDLE_PRIORITY_P2P_WITH_COE
        else:
            rx_tx_pair_id = get_rx_tx_pair_id(ref_test, REF, active_dut, active_ref)
            if rx_tx_pair_id and rx_tx_pair_id not in processed_ref:
                ref_bundled_tests.append(rx_tx_pair_id)
                processed_ref.add(rx_tx_pair_id)
            priority = BUNDLE_PRIORITY_P2P_ONLY

        ref_bundled_tests = dedupe_preserve_order(ref_bundled_tests)
        test_bundles.append(
            create_bundle(ref_bundled_tests, REF, priority, bundle_index, active_dut, active_ref)
        )
        bundle_index += 1

    # Phase 5: unmatched REF P3P bundles
    for ref_test_id, ref_test in active_ref.items():
        if ref_test_id in processed_ref or ref_test.test_type != "P3P":
            continue

        ref_bundled_tests = [ref_test_id]
        processed_ref.add(ref_test_id)

        th_ut_pair_id = get_th_ut_pair_id(ref_test, REF, active_dut, active_ref)
        if th_ut_pair_id and th_ut_pair_id not in processed_ref:
            ref_bundled_tests.append(th_ut_pair_id)
            processed_ref.add(th_ut_pair_id)

        ref_bundled_tests = dedupe_preserve_order(ref_bundled_tests)
        test_bundles.append(
            create_bundle(ref_bundled_tests, REF, BUNDLE_PRIORITY_P3P, bundle_index, active_dut, active_ref)
        )
        bundle_index += 1

    # Phase 6: REF leftovers (including COE-only)
    for ref_test_id in active_ref:
        if ref_test_id in processed_ref:
            continue
        test_bundles.append(
            create_bundle([ref_test_id], REF, BUNDLE_PRIORITY_COE_ONLY, bundle_index, active_dut, active_ref)
        )
        bundle_index += 1
        processed_ref.add(ref_test_id)

    test_bundles.sort(key=lambda b: (b.priority, b.index, b.device, b.config))
    return test_bundles

def get_rx_tx_pair_id(test: Test, device: str, active_dut: dict[str, Test], active_ref: dict[str, Test]) -> str | None:
    active_tests = active_dut if device == DUT else active_ref
    test_id = test.test_id
    pair_id = test_id.replace("RX", "TX", 1) if "RX" in test_id else test_id.replace("TX", "RX", 1)
    return pair_id if pair_id in active_tests else None

def get_th_ut_pair_id(test: Test, device: str, active_dut: dict[str, Test], active_ref: dict[str, Test]) -> str | None:
    active_tests = active_dut if device == DUT else active_ref
    test_id = test.test_id
    pair_id = test_id.replace("TH", "UT", 1) if "TH" in test_id else test_id.replace("UT", "TH", 1)
    return pair_id if pair_id in active_tests else None

def create_bundle(tests: list[str], device: str, priority: int, index: int, active_dut: dict[str, Test], active_ref: dict[str, Test]) -> TestBundle:
    if not tests:
        raise ValueError("Cannot create bundle with no tests")

    if device == DUT:
        active_tests = active_dut
    else:
        active_tests = active_ref

    total_minutes = sum(active_tests[test_id].estimated_minutes for test_id in tests)

    bundle_tests = [active_tests[test_id] for test_id in tests]
    config = tuple(build_bundle_test_configs(bundle_tests))

    return TestBundle(
        index=index,
        tests=tests,
        total_minutes=total_minutes,
        device=device,
        config=config,
        priority=priority,
    )

def bundle_pair_lookup(bundles: list[TestBundle], dut_bundles: list[TestBundle], ref_bundles: list[TestBundle]) -> list[TestBundle]:
    device = bundles[0].device if bundles else None
    active_bundle_ids = {b.index for b in dut_bundles} if device == REF else {b.index for b in ref_bundles}

    bundle_indexes = [b.index for b in bundles if b.index in active_bundle_ids]
    if device == DUT:
        return [b for b in ref_bundles if b.index in bundle_indexes]
    elif device == REF:
        return [b for b in dut_bundles if b.index in bundle_indexes]
    return []