import json import sqlite3 import os import re from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Any, Iterator APP_ROOT = Path(__file__).resolve().parent DB_PATH = APP_ROOT / "scheduler.db" # Hardware device names (from environment or defaults) DUT = os.getenv("DUT", "CGW453").strip() REF = os.getenv("REF", "CGW452").strip() # Device names for the test database (use hardware device names) DEVICE_DUT = DUT DEVICE_REF = REF @dataclass(frozen=True) class TestRecord: test_id: str device: str test_type: str rotation: str | None rx_tx: str | None has_coe_pair: bool coe_pairing: list[str] priority: int victim_band: str | None config: dict[str, dict[str, str | None]] throttled: bool estimated_minutes: int status: str = "pending" excluded: bool = False raw_payload: dict[str, Any] | None = None @dataclass(frozen=True) class ScheduleRow: test_id: str device: str scheduled_date: str shift_index: int sequence_in_shift: int test_type: str rotation: str | None config: dict[str, dict[str, str | None]] status: str priority: int estimated_minutes: int @contextmanager def get_connection(db_path: str | Path = DB_PATH) -> Iterator[sqlite3.Connection]: conn = sqlite3.connect(str(db_path)) conn.row_factory = sqlite3.Row try: yield conn conn.commit() finally: conn.close() def init_db(db_path: str | Path = DB_PATH) -> None: # Validate device names are set if not DUT or not REF: raise ValueError(f"Invalid device names: DUT={DUT!r}, REF={REF!r}") with get_connection(db_path) as conn: conn.executescript( """ CREATE TABLE IF NOT EXISTS tests ( id INTEGER PRIMARY KEY AUTOINCREMENT, test_id TEXT NOT NULL, device TEXT NOT NULL, test_type TEXT NOT NULL CHECK (test_type IN ('P2P', 'COE', 'P3P')), rotation TEXT, rx_tx TEXT CHECK (rx_tx IN ('RX', 'TX') OR rx_tx IS NULL), has_coe_pair INTEGER NOT NULL DEFAULT 0, coe_pairing_json TEXT, priority INTEGER NOT NULL CHECK (priority BETWEEN 1 AND 5), victim_band TEXT, config_json TEXT, throttled INTEGER NOT NULL DEFAULT 0, estimated_minutes INTEGER NOT NULL, status TEXT NOT NULL CHECK (status IN ('pending', 'completed', 'failed', 'invalid')), excluded INTEGER NOT NULL DEFAULT 0, raw_payload TEXT, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, UNIQUE(test_id, device) ); CREATE TABLE IF NOT EXISTS schedules ( id INTEGER PRIMARY KEY AUTOINCREMENT, schedule_version INTEGER NOT NULL, test_id TEXT NOT NULL, device TEXT NOT NULL, scheduled_date TEXT NOT NULL, shift_index INTEGER NOT NULL CHECK (shift_index IN (1, 2, 3)), sequence_in_shift INTEGER NOT NULL, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY(test_id, device) REFERENCES tests(test_id, device) ); CREATE TABLE IF NOT EXISTS settings ( key TEXT PRIMARY KEY, value_json TEXT NOT NULL, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS holidays ( date TEXT PRIMARY KEY, note TEXT ); CREATE TABLE IF NOT EXISTS runtime_defaults ( test_type TEXT PRIMARY KEY, minutes INTEGER NOT NULL ); CREATE TABLE IF NOT EXISTS rerun_events ( id INTEGER PRIMARY KEY AUTOINCREMENT, detected_date TEXT NOT NULL, failed_test_ids_json TEXT NOT NULL, estimated_rerun_minutes INTEGER NOT NULL, rerun_during_day INTEGER ); CREATE INDEX IF NOT EXISTS idx_tests_status ON tests(status); CREATE INDEX IF NOT EXISTS idx_tests_priority ON tests(priority); CREATE INDEX IF NOT EXISTS idx_schedules_date_shift ON schedules(scheduled_date, shift_index); """ ) _ensure_column(conn, "tests", "has_coe_pair", "INTEGER NOT NULL DEFAULT 0") _ensure_column(conn, "tests", "coe_pairing_json", "TEXT") _ensure_column(conn, "tests", "config_json", "TEXT") _ensure_column(conn, "tests", "victim_band", "TEXT") _ensure_column(conn, "tests", "excluded", "INTEGER NOT NULL DEFAULT 0") _ensure_column(conn, "tests", "throttled", "INTEGER NOT NULL DEFAULT 0") # Seed runtime defaults for schedule estimation. conn.execute( """ INSERT INTO runtime_defaults(test_type, minutes) VALUES ('P2P', 80), ('COE', 115), ('P3P', 105) ON CONFLICT(test_type) DO NOTHING """ ) def upsert_tests(records: list[TestRecord], db_path: str | Path = DB_PATH) -> int: if not records: return 0 with get_connection(db_path) as conn: conn.executemany( """ INSERT INTO tests( test_id, device, test_type, rotation, rx_tx, has_coe_pair, coe_pairing_json, config_json, priority, victim_band, throttled, estimated_minutes, status, excluded, raw_payload ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(test_id, device) DO UPDATE SET test_type = excluded.test_type, device = excluded.device, rotation = excluded.rotation, rx_tx = excluded.rx_tx, has_coe_pair = excluded.has_coe_pair, coe_pairing_json = excluded.coe_pairing_json, config_json = excluded.config_json, priority = excluded.priority, victim_band = excluded.victim_band, throttled = excluded.throttled, estimated_minutes = excluded.estimated_minutes, status = excluded.status, excluded = excluded.excluded, raw_payload = excluded.raw_payload, updated_at = CURRENT_TIMESTAMP """, [ ( r.test_id, r.device, r.test_type, r.rotation, r.rx_tx, int(r.has_coe_pair), json.dumps(r.coe_pairing or []), json.dumps(r.config), r.priority, r.victim_band, int(r.throttled), r.estimated_minutes, r.status, int(r.excluded), json.dumps(r.raw_payload or {}), ) for r in records ], ) return len(records) def _ensure_column(conn: sqlite3.Connection, table: str, column: str, column_type: str) -> None: existing = conn.execute(f"PRAGMA table_info({table})").fetchall() column_names = {row[1] for row in existing} if column not in column_names: conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {column_type}") def save_settings(payload: dict[str, Any], db_path: str | Path = DB_PATH) -> None: with get_connection(db_path) as conn: for key, value in payload.items(): conn.execute( """ INSERT INTO settings(key, value_json, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP) ON CONFLICT(key) DO UPDATE SET value_json = excluded.value_json, updated_at = CURRENT_TIMESTAMP """, (key, json.dumps(value)), ) def read_settings(db_path: str | Path = DB_PATH) -> dict[str, Any]: with get_connection(db_path) as conn: rows = conn.execute("SELECT key, value_json FROM settings").fetchall() return {row["key"]: json.loads(row["value_json"]) for row in rows} def _parse_rule_tokens(rule: str | None) -> list[str]: if not rule: return [] return [token.strip().upper() for token in rule.split(",") if token.strip()] def _normalize_text(value: Any) -> str: if value is None: return "" return " ".join(str(value).strip().upper().split()) def _normalize_compact(value: Any) -> str: return re.sub(r"[^A-Z0-9+\-]", "", _normalize_text(value)) def _extract_signed_int(value: Any) -> str: match = re.search(r"[-+]?\d+", _normalize_text(value)) return match.group(0) if match else "" def _band_entry(config: dict[str, dict[str, str | None]], band: str) -> dict[str, str | None]: return config.get(band) or config.get(band.lower()) or {} def _entry_value(entry: dict[str, str | None], key: str) -> str: if not entry: return "" return _normalize_text(entry.get(key) or entry.get(key.capitalize()) or entry.get(key.upper())) def _record_power_mode(record: TestRecord) -> str: power_mode = _entry_value(_band_entry(record.config, "6G"), "power_mode") if power_mode: return power_mode raw_payload = record.raw_payload or {} return _normalize_text(raw_payload.get("6GHz Power Mode")) def _match_any_band_value(record: TestRecord, key: str, token_suffix: str) -> bool: if not token_suffix: return False target_num = _extract_signed_int(token_suffix) target_compact = _normalize_compact(token_suffix) keys = ("Station 1", "Station 2", "Station 3") if record.test_type == "P3P" else ("5G", "6G", "2G") for band in keys: entry = _band_entry(record.config, band) value = _entry_value(entry, key) if not value: continue if target_num: if _extract_signed_int(value) == target_num: return True elif target_compact and _normalize_compact(value) == target_compact: return True return False def _matches_atomic_exclusion_token(record: TestRecord, token: str) -> bool: if not token: return False device = _normalize_text(record.device) test_type = _normalize_text(record.test_type) rotation = _normalize_text(record.rotation) victim_band = _normalize_text(record.victim_band) power_mode = _record_power_mode(record) known_devices = {_normalize_text(DUT), _normalize_text(REF)} if token in {"LPI", "SP"}: return power_mode == token if token in known_devices: return device == token if token in {"P2P", "P3P", "COE"}: return test_type == token if token in {"2G", "5G", "6G"}: return victim_band == token if token.startswith("BW"): return _match_any_band_value(record, "bandwidth", token[2:]) if token.startswith("CH"): return _match_any_band_value(record, "channel", token[2:]) if token.startswith("RSSI"): return _match_any_band_value(record, "rssi", token[4:]) if token.startswith("R") and len(token) > 1: return rotation == token return False def _matches_exclusion_token(record: TestRecord, token: str) -> bool: if not token: return False parts = [part for part in token.split("_") if part] if len(parts) > 1: # Treat underscore as logical AND across sub-tokens. return all(_matches_atomic_exclusion_token(record, part) for part in parts) return _matches_atomic_exclusion_token(record, token) def _should_exclude_by_rule(record: TestRecord, tokens: list[str]) -> bool: return any(_matches_exclusion_token(record, token) for token in tokens) def list_schedulable_tests(db_path: str | Path = DB_PATH, rule: str = "") -> list[TestRecord]: tokens = _parse_rule_tokens(rule) with get_connection(db_path) as conn: rows = conn.execute( """ SELECT test_id, device, test_type, rotation, rx_tx, has_coe_pair, config_json, victim_band, coe_pairing_json, priority, throttled, estimated_minutes, excluded, status, raw_payload FROM tests WHERE excluded = 0 AND status != 'completed' """ ).fetchall() results: list[TestRecord] = [] for row in rows: results.append( TestRecord( test_id=row["test_id"], device=row["device"], test_type=row["test_type"], rotation=row["rotation"], rx_tx=row["rx_tx"], has_coe_pair=bool(row["has_coe_pair"]), coe_pairing=json.loads(row["coe_pairing_json"] or "[]"), priority=int(row["priority"]), victim_band=row["victim_band"], config=json.loads(row["config_json"] or "{}"), throttled=bool(row["throttled"]), estimated_minutes=int(row["estimated_minutes"]), status=row["status"], excluded=bool(row["excluded"]), raw_payload=json.loads(row["raw_payload"] or "{}"), ) ) if not tokens: return results return [record for record in results if not _should_exclude_by_rule(record, tokens)] # Get all tests for a specific device, regardless of exclusion or completion status. Used for schedule display and management. def list_tests_for_device(device: str, db_path: str | Path = DB_PATH) -> list[TestRecord]: with get_connection(db_path) as conn: rows = conn.execute( """ SELECT test_id, device, test_type, rotation, rx_tx, has_coe_pair, config_json, victim_band, coe_pairing_json, priority, throttled, estimated_minutes, excluded, status, raw_payload FROM tests WHERE device = ? """, (device,), ).fetchall() results: list[TestRecord] = [] for row in rows: results.append( TestRecord( test_id=row["test_id"], device=row["device"], test_type=row["test_type"], rotation=row["rotation"], rx_tx=row["rx_tx"], has_coe_pair=bool(row["has_coe_pair"]), coe_pairing=json.loads(row["coe_pairing_json"] or "[]"), priority=int(row["priority"]), victim_band=row["victim_band"], config=json.loads(row["config_json"] or "{}"), throttled=bool(row["throttled"]), estimated_minutes=int(row["estimated_minutes"]), status=row["status"], excluded=bool(row["excluded"]), raw_payload=json.loads(row["raw_payload"] or "{}"), ) ) return results def list_holidays(db_path: str | Path = DB_PATH) -> set[str]: with get_connection(db_path) as conn: rows = conn.execute("SELECT date FROM holidays").fetchall() return {row["date"] for row in rows} def create_schedule_version( entries: list[tuple[str, str, str, int, int]], db_path: str | Path = DB_PATH, ) -> int: with get_connection(db_path) as conn: row = conn.execute("SELECT COALESCE(MAX(schedule_version), 0) AS current FROM schedules").fetchone() next_version = int(row["current"]) + 1 if entries: conn.executemany( """ INSERT INTO schedules(schedule_version, test_id, device, scheduled_date, shift_index, sequence_in_shift) VALUES (?, ?, ?, ?, ?, ?) """, [ (next_version, test_id, device, scheduled_date, shift_index, sequence) for test_id, device, scheduled_date, shift_index, sequence in entries ], ) return next_version def get_schedule_week(start_date: str, db_path: str | Path = DB_PATH) -> list[ScheduleRow]: with get_connection(db_path) as conn: version_row = conn.execute("SELECT MAX(schedule_version) AS latest FROM schedules").fetchone() latest = version_row["latest"] if latest is None: return [] rows = conn.execute( """ SELECT s.test_id, s.device, s.scheduled_date, s.shift_index, s.sequence_in_shift, t.test_type, t.rotation, t.config_json, t.status, t.priority, t.estimated_minutes FROM schedules s JOIN tests t ON t.test_id = s.test_id AND t.device = s.device WHERE s.schedule_version = ? AND s.scheduled_date >= ? AND s.scheduled_date < date(?, '+7 day') ORDER BY s.scheduled_date, s.shift_index, s.sequence_in_shift """, (latest, start_date, start_date), ).fetchall() result: list[ScheduleRow] = [] for row in rows: result.append( ScheduleRow( test_id=row["test_id"], device=row["device"], scheduled_date=row["scheduled_date"], shift_index=int(row["shift_index"]), sequence_in_shift=int(row["sequence_in_shift"]), test_type=row["test_type"], rotation=row["rotation"], config=json.loads(row["config_json"] or "{}"), status=row["status"], priority=int(row["priority"]), estimated_minutes=int(row["estimated_minutes"]), ) ) return result def mark_tests_completed(test_ids_with_device: list[tuple[str, str]], db_path: str | Path = DB_PATH) -> None: if not test_ids_with_device: return with get_connection(db_path) as conn: conn.executemany( """ UPDATE tests SET status = 'completed', updated_at = CURRENT_TIMESTAMP WHERE test_id = ? AND device = ? AND status != 'completed' """, test_ids_with_device, ) def upsert_holidays(dates: list[str], db_path: str | Path = DB_PATH) -> None: """Replace all holidays with the provided list of YYYY-MM-DD date strings.""" with get_connection(db_path) as conn: conn.execute("DELETE FROM holidays") if dates: conn.executemany( "INSERT OR IGNORE INTO holidays(date) VALUES (?)", [(d,) for d in dates], )