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 from test_config import serialize_station_testpoint_map APP_ROOT = Path(__file__).resolve().parent DB_PATH = Path(os.getenv("DB_PATH", str(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 power_mode: 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 station_testpoint_map: str | 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), power_mode TEXT, 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', 'rerun')), 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 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", "power_mode", "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") _ensure_column(conn, "tests", "station_testpoint_map", "TEXT") # 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 _serialize_station_testpoint_map(config: dict[str, dict[str, str | None]]) -> str: """Compute and serialize the station-to-testpoint map from a test config.""" return serialize_station_testpoint_map(config) def upsert_tests(records: list[TestRecord], db_path: str | Path = DB_PATH) -> int: if not records: return 0 values: list[tuple[Any, ...]] = [] for r in records: station_testpoint_map = _serialize_station_testpoint_map(r.config) values.append( ( r.test_id, r.device, r.test_type, r.rotation, r.rx_tx, r.power_mode, 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 {}), station_testpoint_map, ) ) with get_connection(db_path) as conn: conn.executemany( """ INSERT INTO tests( test_id, device, test_type, rotation, rx_tx, power_mode, has_coe_pair, coe_pairing_json, config_json, priority, victim_band, throttled, estimated_minutes, status, excluded, raw_payload, station_testpoint_map ) 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, power_mode = excluded.power_mode, 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 = CASE WHEN tests.status IN ('completed', 'rerun') THEN tests.status ELSE excluded.status END, excluded = excluded.excluded, raw_payload = excluded.raw_payload, station_testpoint_map = excluded.station_testpoint_map, updated_at = CURRENT_TIMESTAMP """, values, ) 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: return _normalize_text(record.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, power_mode, 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"], power_mode=row["power_mode"], 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, power_mode, config_json, victim_band, coe_pairing_json, priority, throttled, estimated_minutes, excluded, status, raw_payload, station_testpoint_map 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"], power_mode=row["power_mode"], 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 "{}"), station_testpoint_map=row["station_testpoint_map"], ) ) 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_versions(db_path: str | Path = DB_PATH) -> list[dict[str, Any]]: with get_connection(db_path) as conn: rows = conn.execute( """ SELECT schedule_version, MIN(created_at) AS created_at, COUNT(*) AS entry_count FROM schedules GROUP BY schedule_version ORDER BY schedule_version DESC """ ).fetchall() return [ { "schedule_version": int(row["schedule_version"]), "created_at": row["created_at"], "entry_count": int(row["entry_count"]), } for row in rows ] def resolve_schedule_version( version: int | None = None, db_path: str | Path = DB_PATH, ) -> int | None: with get_connection(db_path) as conn: if version is not None: row = conn.execute( "SELECT 1 FROM schedules WHERE schedule_version = ? LIMIT 1", (version,), ).fetchone() return int(version) if row else None latest_row = conn.execute("SELECT MAX(schedule_version) AS latest FROM schedules").fetchone() latest = latest_row["latest"] return int(latest) if latest is not None else None def get_schedule_week( start_date: str, version: int | None = None, db_path: str | Path = DB_PATH, ) -> list[ScheduleRow]: with get_connection(db_path) as conn: if version is not None: version_row = conn.execute( "SELECT 1 AS exists_row FROM schedules WHERE schedule_version = ? LIMIT 1", (version,), ).fetchone() if version_row is None: return [] selected_version = int(version) else: latest_row = conn.execute("SELECT MAX(schedule_version) AS latest FROM schedules").fetchone() latest = latest_row["latest"] if latest is None: return [] selected_version = int(latest) if selected_version 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') AND s.shift_index IN (2, 3) ) OR ( s.scheduled_date > ? AND s.scheduled_date < date(?, '+7 day') AND s.shift_index = 1 ) OR ( s.scheduled_date = date(?, '+7 day') AND s.shift_index = 1 ) ) ORDER BY s.scheduled_date, s.shift_index, s.sequence_in_shift """, ( selected_version, start_date, start_date, start_date, 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 get_schedule_rows( version: int | None = None, db_path: str | Path = DB_PATH, ) -> list[ScheduleRow]: with get_connection(db_path) as conn: if version is not None: version_row = conn.execute( "SELECT 1 AS exists_row FROM schedules WHERE schedule_version = ? LIMIT 1", (version,), ).fetchone() if version_row is None: return [] selected_version = int(version) else: latest_row = conn.execute("SELECT MAX(schedule_version) AS latest FROM schedules").fetchone() latest = latest_row["latest"] if latest is None: return [] selected_version = int(latest) if selected_version 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 = ? ORDER BY s.scheduled_date, s.shift_index, s.sequence_in_shift """, (selected_version,), ).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 reset_completed_to_pending(db_path: str | Path = DB_PATH) -> int: """Reset all 'completed' tests back to 'pending'. Used before a full directory rescan.""" with get_connection(db_path) as conn: cursor = conn.execute( """ UPDATE tests SET status = 'pending', updated_at = CURRENT_TIMESTAMP WHERE status = 'completed' """ ) return cursor.rowcount def mark_tests_completed(test_ids_with_device: list[tuple[str, str]], db_path: str | Path = DB_PATH) -> int: if not test_ids_with_device: return 0 with get_connection(db_path) as conn: # executemany() doesn't properly return rowcount, so we track manually total_updated = 0 for test_id, device in test_ids_with_device: cursor = conn.execute( """ UPDATE tests SET status = 'completed', updated_at = CURRENT_TIMESTAMP WHERE test_id = ? AND device = ? AND status != 'completed' """, (test_id, device), ) if cursor.rowcount > 0: total_updated += cursor.rowcount print(f"[db] marked {test_id} on {device} as completed") else: print(f"[db] {test_id} on {device}: no update (not found or already completed)") return total_updated def mark_overdue_as_rerun(db_path: str | Path = DB_PATH) -> int: """Mark tests from the last overnight window that are still pending as rerun. 'Last overnight window' = all pending tests scheduled before today (any shift) plus today's shift 1 (1am–10am) if it has already ended (current hour >= 10). Returns the number of tests newly marked as rerun. """ from datetime import date as _date, datetime as _datetime now = _datetime.now() today = now.date().isoformat() shift1_ended = now.hour >= 10 # shift 1 ends ~10am 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 0 if shift1_ended: rows = conn.execute( """ SELECT DISTINCT s.test_id, s.device FROM schedules s JOIN tests t ON t.test_id = s.test_id AND t.device = s.device WHERE s.schedule_version = ? AND t.status = 'pending' AND ( s.scheduled_date < ? OR (s.scheduled_date = ? AND s.shift_index = 1) ) """, (latest, today, today), ).fetchall() else: rows = conn.execute( """ SELECT DISTINCT s.test_id, s.device 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 t.status = 'pending' """, (latest, today), ).fetchall() if not rows: return 0 conn.executemany( """ UPDATE tests SET status = 'rerun', updated_at = CURRENT_TIMESTAMP WHERE test_id = ? AND device = ? AND status = 'pending' """, [(r["test_id"], r["device"]) for r in rows], ) return len(rows) def get_rerun_tests(db_path: str | Path = DB_PATH) -> list[dict]: """Return tests scheduled before shift 2 today that are not yet completed. This checks the latest schedule version and returns tests from: - Yesterday's shift 3 (5pm-1am) - Today's shift 1 (1am-10am) Once tests are rescheduled to shift 2 or later today, they no longer appear. """ from datetime import date as _date today = _date.today().isoformat() 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 DISTINCT t.test_id, t.device, t.test_type, t.estimated_minutes FROM tests t JOIN schedules s ON s.test_id = t.test_id AND s.device = t.device WHERE s.schedule_version = ? AND t.status != 'completed' AND ( (s.scheduled_date = date(?, '-1 day') AND s.shift_index = 3) OR (s.scheduled_date = ? AND s.shift_index = 1) ) ORDER BY t.test_id, t.device """, (latest, today, today), ).fetchall() return [ { "test_id": r["test_id"], "device": r["device"], "test_type": r["test_type"], "estimated_minutes": r["estimated_minutes"], } for r in rows ] 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], ) def reset_tests(db_path: str | Path = DB_PATH) -> int: """Clear all tests records keep empty tests table.""" with get_connection(db_path) as conn: cursor = conn.execute("DELETE FROM tests") return cursor.rowcount