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=False)
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
    status: str

# Bundle priority tiers for scheduling order (lower number = higher priority)
BUNDLE_PRIORITY_DUT_COMPLETED_MIRROR = 0  # REF-side tests whose DUT counterpart is already completed
BUNDLE_PRIORITY_USER_TOP = 1  # User-specified top priority tests
BUNDLE_PRIORITY_RERUN = 2   # Rerun tests (DUT or REF) 
BUNDLE_PRIORITY_MIRROR = 3
BUNDLE_PRIORITY_P2P_WITH_COE = 4  # P2P tests with COE pairs
BUNDLE_PRIORITY_P2P_ONLY = 5  # P2P tests without COE pairs (RX/TX bundled)
BUNDLE_PRIORITY_COE_ONLY = 6  # COE tests without P2P pairing (should be rare/unschedulable)
BUNDLE_PRIORITY_P3P = 6  # P3P tests

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

@dataclass(frozen=False)
class TestBundle:
    index: int
    tests: list[str]
    total_minutes: int
    device: str
    config: tuple[str, ...]
    priority: int
    completed: 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."""

    print(f"Top priority tests: {top_priority_tests}")
    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

    # 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)

        # Check if any of the DUT tests in this bundle are in the top priority list
        if any((dut_test_id, DUT) in top_priority_tests for dut_test_id in dut_bundled_tests):
            priority = BUNDLE_PRIORITY_USER_TOP

        # Check if any of the DUT tests in this bundle are reruns
        elif any((dut_test.status == "rerun" for dut_test_id in dut_bundled_tests if (dut_test := active_dut.get(dut_test_id)))):
            priority = BUNDLE_PRIORITY_RERUN
        
        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)

            # Check if any of the REF tests in this bundle are in the top priority list
            if any((ref_test_id, REF) in top_priority_tests for ref_test_id in ref_bundled_tests):
                priority = BUNDLE_PRIORITY_USER_TOP
            # Check if any of the REF tests in this bundle are reruns
            elif any((ref_test.status == "rerun" for ref_test_id in ref_bundled_tests if (ref_test := active_ref.get(ref_test_id)))):
                priority = BUNDLE_PRIORITY_RERUN

            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)
        priority = BUNDLE_PRIORITY_P3P

        # Check if any of the DUT tests in this bundle are in the top priority list
        if any((dut_test_id, DUT) in top_priority_tests for dut_test_id in dut_bundled_tests):
            priority = BUNDLE_PRIORITY_USER_TOP

        # Check if any of the DUT tests in this bundle are reruns
        elif any((dut_test.status == "rerun" for dut_test_id in dut_bundled_tests if (dut_test := active_dut.get(dut_test_id)))):
            priority = BUNDLE_PRIORITY_RERUN
            
        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
        ]

        # Check if any of the REF tests in this bundle are in the top priority list
        if any((ref_test_id, REF) in top_priority_tests for ref_test_id in ref_bundled_tests):
            priority = BUNDLE_PRIORITY_USER_TOP
        # Check if any of the REF tests in this bundle are reruns
        elif any((ref_test.status == "rerun" for ref_test_id in ref_bundled_tests if (ref_test := active_ref.get(ref_test_id)))):
            priority = BUNDLE_PRIORITY_RERUN

        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 3: DUT leftovers (COE-only)
    for dut_test_id in active_dut:
        if dut_test_id in processed_dut:
            continue
        print(f"Orphan COE-only DUT test found: {dut_test_id}")
        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 bundles, meaning DUT is completed
    for ref_test_id, ref_test in active_ref.items():
        if ref_test_id in processed_ref:
            continue

        ref_bundled_tests = [ref_test_id]
        processed_ref.add(ref_test_id)

        if ref_test.test_type == "P2P":
            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_DUT_COMPLETED_MIRROR
            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_DUT_COMPLETED_MIRROR

        if ref_test.test_type == "P3P":
            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)
            priority = BUNDLE_PRIORITY_DUT_COMPLETED_MIRROR

        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: REF leftovers ( COE-only)
    for ref_test_id in active_ref:
        if ref_test_id in processed_ref:
            continue
        print(f"Orphan COE-only REF test found: {ref_test_id}")
        test_bundles.append(
            create_bundle([ref_test_id], REF, BUNDLE_PRIORITY_DUT_COMPLETED_MIRROR, 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,
        completed=0
    )

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 []

def update_mirrored_bundle_priorities(knapsack_bundles: list[TestBundle], dut_pending: list[TestBundle], ref_pending: list[TestBundle]) -> tuple[list[TestBundle], list[TestBundle]]:
    bundle_pairs = bundle_pair_lookup(knapsack_bundles, dut_pending, ref_pending)

    for bundle in dut_pending:
        if bundle.index in {b.index for b in bundle_pairs}:
            bundle.priority = BUNDLE_PRIORITY_MIRROR

    for bundle in ref_pending:
        if bundle.index in {b.index for b in bundle_pairs}:
            bundle.priority = BUNDLE_PRIORITY_MIRROR

    return dut_pending, ref_pending