from contextlib import asynccontextmanager from pathlib import Path from typing import Any from datetime import date, datetime, timedelta import os from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field import db from parser import CsvValidationError, parse_target_csv # Ensure new_scheduler resolves the same DUT/REF labels as db records. os.environ.setdefault("DUT", db.DEVICE_DUT) os.environ.setdefault("REF", db.DEVICE_REF) from scheduler import Scheduler, Test as SchedulerTest from scanner import resolve_runtime_path from test_config import build_bundle_test_configs, build_config_rows from test_window import get_shift_capacity_for_date, get_shift_sequence, is_off_day from watcher import configure_result_watcher, stop_result_watcher APP_ROOT = Path(__file__).resolve().parent DB_PATH = Path(os.getenv("DB_PATH", str(APP_ROOT / "scheduler.db"))) DUT = os.getenv("DUT", "CGW453").strip() REF = os.getenv("REF", "CGW452").strip() class LoadTestsRequest(BaseModel): csv_path: str = Field(..., description="Absolute or backend-relative path to target CSV") class SaveSettingsRequest(BaseModel): settings: dict[str, Any] def _smb_credentials_from_settings(settings: dict[str, Any]) -> dict[str, str]: return { "username": str(settings.get("smbUsername") or "").strip(), "password": str(settings.get("smbPassword") or ""), "domain": str(settings.get("smbDomain") or "").strip(), } def _runtime_overrides_from_settings(settings: dict[str, Any]) -> dict[str, dict[str, int]]: def _parse_positive_int(value: Any) -> int | None: if value is None: return None if isinstance(value, str): value = value.strip() if not value: return None try: parsed = int(value) except (TypeError, ValueError): return None return parsed if parsed > 0 else None legacy_overrides = { "P2P": _parse_positive_int(settings.get("p2pRuntimeMinutes")), "COE": _parse_positive_int(settings.get("coeRuntimeMinutes")), "P3P": _parse_positive_int(settings.get("p3pRuntimeMinutes")), } overrides: dict[str, dict[str, int]] = {} for device_key, field_prefix in ((db.DEVICE_DUT, "dut"), (db.DEVICE_REF, "ref")): device_overrides: dict[str, int] = {} for test_type, suffix in (("P2P", "P2p"), ("COE", "Coe"), ("P3P", "P3p")): minutes = _parse_positive_int(settings.get(f"{field_prefix}{suffix}RuntimeMinutes")) if minutes is None: minutes = legacy_overrides[test_type] if minutes is not None: device_overrides[test_type] = minutes if device_overrides: overrides[device_key] = device_overrides return overrides def _is_explicit_path(path_value: str) -> bool: return ( path_value.startswith(("/", "\\\\", "//")) or "\\" in path_value or ":" in path_value ) def _resolve_requested_csv_path(requested_path: str) -> str | Path: received_path = str(requested_path).strip() resolved_path = str(resolve_runtime_path(received_path)).strip() if _is_explicit_path(received_path) or _is_explicit_path(resolved_path): return resolved_path data_dir_path = APP_ROOT / "data" / received_path if data_dir_path.exists(): return data_dir_path return APP_ROOT / received_path class CompileScheduleRequest(BaseModel): start_date: str | None = Field(default=None, description="YYYY-MM-DD") rule: str = "" daytime_testing_today: bool = False daytime_testing_hours: int = Field(default=8, ge=0, le=8) daytime_testing_device: str = "" dual_device_weekend_start_enabled: bool = False dual_device_weekend_start_dates: list[str] = Field(default_factory=list) top_priority_tests_dut: list[str] = Field(default_factory=list) top_priority_tests_ref: list[str] = Field(default_factory=list) top_priority_tests: list[str] = Field(default_factory=list) lowest_priority_tests: list[str] = Field(default_factory=list) class RemoveActiveTestsRequest(BaseModel): test_ids: list[str] = Field(default_factory=list) class SaveHolidaysRequest(BaseModel): dates: list[str] = Field(default_factory=list) SHIFT_LABELS = { 1: "12AM–9AM", 2: "9AM–5PM", 3: "5PM–12AM", } SHIFT_BOUNDARIES = { 1: ((0, 0), (9, 0)), 2: ((9, 0), (17, 0)), 3: ((17, 0), (24, 0)), } def _serialize_schedule_row(row: db.ScheduleRow) -> dict[str, Any]: return { "test_id": row.test_id, "device": row.device, "scheduled_date": row.scheduled_date, "shift_index": row.shift_index, "sequence_in_shift": row.sequence_in_shift, "test_type": row.test_type, "rotation": row.rotation, "config": row.config, "status": row.status, "priority": row.priority, "estimated_minutes": row.estimated_minutes, } def _window_segment_key(day: date, shift_index: int) -> str: return f"{day.isoformat()}::shift{shift_index}" def _combine_date_time(day: date, hour: int, minute: int) -> str: if hour == 24: return datetime.combine(day + timedelta(days=1), datetime.min.time()).isoformat(timespec="minutes") return datetime.combine(day, datetime.min.time()).replace(hour=hour, minute=minute).isoformat(timespec="minutes") def _window_intersects_week(segments: list[tuple[date, int]], week_start: date, week_end: date) -> bool: for day, _shift_index in segments: if week_start <= day < week_end: return True return False def _build_schedule_windows(week_start: date, all_rows: list[db.ScheduleRow], holiday_dates: set[str]) -> list[dict[str, Any]]: if not all_rows: return [] all_dates = [datetime.strptime(row.scheduled_date, "%Y-%m-%d").date() for row in all_rows] first_scheduled_date = min(all_dates) last_scheduled_date = max(all_dates) week_end = week_start + timedelta(days=7) rows_by_segment: dict[tuple[str, int], list[dict[str, Any]]] = {} weekday_shift2_dates: set[str] = set() for row in all_rows: serialized = _serialize_schedule_row(row) rows_by_segment.setdefault((row.scheduled_date, row.shift_index), []).append(serialized) row_date = datetime.strptime(row.scheduled_date, "%Y-%m-%d").date() if row.shift_index == 2 and not is_off_day(row_date, holiday_dates): weekday_shift2_dates.add(row.scheduled_date) simulation_start = min(first_scheduled_date, week_start) - timedelta(days=7) while is_off_day(simulation_start, holiday_dates): simulation_start -= timedelta(days=1) simulation_end = max(last_scheduled_date, week_end) + timedelta(days=7) current_date = simulation_start windows: list[dict[str, Any]] = [] seen_window_ids: set[str] = set() while current_date <= simulation_end: current_iso = current_date.isoformat() if current_iso in weekday_shift2_dates: segments = get_shift_sequence(current_date, holiday_dates, daytime_shift2_only=True) window_id = f"{current_iso}-shift2" if window_id not in seen_window_ids and _window_intersects_week(segments, week_start, week_end): seen_window_ids.add(window_id) shift_tests = [ test for day, shift_index in segments for test in rows_by_segment.get((day.isoformat(), shift_index), []) ] available_runtime_minutes = sum( get_shift_capacity_for_date(day, holiday_dates, daytime_shift2_only=True).get(shift_index, 0) for day, shift_index in segments ) start_hour, start_minute = SHIFT_BOUNDARIES[segments[0][1]][0] end_hour, end_minute = SHIFT_BOUNDARIES[segments[-1][1]][1] windows.append({ "window_id": window_id, "window_type": "daytime", "start_date": segments[0][0].isoformat(), "start_shift_index": segments[0][1], "end_date": segments[-1][0].isoformat(), "end_shift_index": segments[-1][1], "start_at": _combine_date_time(segments[0][0], start_hour, start_minute), "end_at": _combine_date_time(segments[-1][0], end_hour, end_minute), "available_runtime_minutes": available_runtime_minutes, "estimated_runtime_minutes": sum(test["estimated_minutes"] for test in shift_tests), "bundle_test_configs": build_bundle_test_configs(shift_tests), "segments": [ { "date": day.isoformat(), "shift_index": shift_index, "label": SHIFT_LABELS[shift_index], "segment_key": _window_segment_key(day, shift_index), } for day, shift_index in segments ], "configRows": build_config_rows(shift_tests), "tests": shift_tests, }) segments = get_shift_sequence(current_date, holiday_dates, daytime_shift2_only=False) last_sequence_date = segments[-1][0] last_shift_index = segments[-1][1] window_kind = "offday" if is_off_day(current_date, holiday_dates) else "overnight" window_id = f"{current_iso}-shift{segments[0][1]}" if window_id not in seen_window_ids and _window_intersects_week(segments, week_start, week_end): seen_window_ids.add(window_id) window_tests = [ test for day, shift_index in segments for test in rows_by_segment.get((day.isoformat(), shift_index), []) ] available_runtime_minutes = sum( get_shift_capacity_for_date(day, holiday_dates, daytime_shift2_only=False).get(shift_index, 0) for day, shift_index in segments ) start_hour, start_minute = SHIFT_BOUNDARIES[segments[0][1]][0] end_hour, end_minute = SHIFT_BOUNDARIES[segments[-1][1]][1] windows.append({ "window_id": window_id, "window_type": window_kind, "start_date": segments[0][0].isoformat(), "start_shift_index": segments[0][1], "end_date": segments[-1][0].isoformat(), "end_shift_index": segments[-1][1], "start_at": _combine_date_time(segments[0][0], start_hour, start_minute), "end_at": _combine_date_time(segments[-1][0], end_hour, end_minute), "available_runtime_minutes": available_runtime_minutes, "estimated_runtime_minutes": sum(test["estimated_minutes"] for test in window_tests), "bundle_test_configs": build_bundle_test_configs(window_tests), "segments": [ { "date": day.isoformat(), "shift_index": shift_index, "label": SHIFT_LABELS[shift_index], "segment_key": _window_segment_key(day, shift_index), } for day, shift_index in segments ], "configRows": build_config_rows(window_tests), "tests": window_tests, }) if last_shift_index in {1, 2}: current_date = last_sequence_date else: current_date = last_sequence_date + timedelta(days=1) windows.sort(key=lambda window: (window["start_at"], window["window_id"])) return windows @asynccontextmanager async def lifespan(application: FastAPI): db.init_db(DB_PATH) settings = db.read_settings(DB_PATH) configure_result_watcher(settings) try: yield finally: stop_result_watcher() app = FastAPI(title="Scheduler API", version="0.1.0", lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=[ "http://localhost:5173", "http://127.0.0.1:5173", "http://localhost:8080", "http://127.0.0.1:8080", ], allow_methods=["*"], allow_headers=["*"], ) @app.get("/api/health") def health() -> dict[str, str]: return {"status": "ok"} @app.post("/api/settings") def save_settings(request: SaveSettingsRequest) -> dict[str, str]: db.save_settings(request.settings, DB_PATH) configure_result_watcher(db.read_settings(DB_PATH)) return {"status": "saved"} @app.get("/api/settings") def get_settings() -> dict[str, Any]: return db.read_settings(DB_PATH) @app.post("/api/settings/restart") def restart_and_clear_data() -> dict[str, Any]: deleted_counts = db.reset_all_data(DB_PATH) configure_result_watcher({}) return { "status": "cleared", **deleted_counts, } @app.post("/api/tests/load") def load_tests(request: LoadTestsRequest) -> dict[str, Any]: csv_path = _resolve_requested_csv_path(request.csv_path) settings = db.read_settings(DB_PATH) smb_credentials = _smb_credentials_from_settings(settings) runtime_overrides = _runtime_overrides_from_settings(settings) try: parsed = parse_target_csv( csv_path, smb_credentials=smb_credentials, runtime_overrides=runtime_overrides, ) except FileNotFoundError as exc: raise HTTPException(status_code=404, detail=str(exc)) from exc except CsvValidationError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc count = db.upsert_tests(parsed.tests, DB_PATH) print(f"Loaded {count} tests from {csv_path}, with {len(parsed.warnings)} warnings.") return { "loaded_tests": count, "warnings": parsed.warnings, } @app.post("/api/schedule/active/remove") def remove_active_tests(request: RemoveActiveTestsRequest) -> dict[str, Any]: # new_scheduler does not keep global mutable active state in app lifecycle. # Keep endpoint for compatibility with frontend calls. return {"status": "ok", "removed": len(request.test_ids)} @app.post("/api/schedule/compile") def compile_schedule_endpoint(request: CompileScheduleRequest) -> dict[str, Any]: if request.start_date: try: datetime.strptime(request.start_date, "%Y-%m-%d") except ValueError as exc: raise HTTPException(status_code=400, detail="start_date must be YYYY-MM-DD") from exc requested_daytime_device = str(request.daytime_testing_device or "").strip() if not requested_daytime_device: requested_daytime_device = DUT if requested_daytime_device not in {DUT, REF}: raise HTTPException(status_code=400, detail=f"daytime_testing_device must be {DUT} or {REF}") all_tests = db.get_not_excluded_tests(DB_PATH, rule=request.rule) stored_tests = db.list_schedulable_tests(DB_PATH, rule=request.rule) if not stored_tests: version = db.create_schedule_version([], DB_PATH) return { "schedule_version": version, "scheduled_tests": 0, "completion_date": None, } scheduler_tests = [ SchedulerTest( test_id=t.test_id, device=t.device, test_type=t.test_type, rotation=t.rotation, rx_tx=t.rx_tx, has_coe_pair=t.has_coe_pair, coe_pairing=t.coe_pairing or [], config=t.config, throttled=t.throttled, estimated_minutes=t.estimated_minutes, status=t.status, ) for t in stored_tests ] holiday_dates = db.list_holidays(DB_PATH) top_priority_pairs: set[tuple[str, str]] = set() dut_top_priority_ids = {item.strip() for item in request.top_priority_tests_dut if item.strip()} ref_top_priority_ids = {item.strip() for item in request.top_priority_tests_ref if item.strip()} status_lookup = {(test.test_id, test.device): test.status for test in stored_tests} for test_id in dut_top_priority_ids: if any(t.test_id == test_id and t.device == DUT for t in stored_tests): top_priority_pairs.add((test_id, DUT)) for test_id in ref_top_priority_ids: if any(t.test_id == test_id and t.device == REF for t in stored_tests): top_priority_pairs.add((test_id, REF)) scheduler = Scheduler( all_tests=all_tests, schedulable_tests=scheduler_tests, top_priority_tests=top_priority_pairs, start_date=request.start_date, holiday_dates=holiday_dates, daytime_testing_today=request.daytime_testing_today, daytime_testing_hours=request.daytime_testing_hours, daytime_testing_device=requested_daytime_device, dual_device_weekend_start_enabled=request.dual_device_weekend_start_enabled, dual_device_window_start_dates=set(item.strip() for item in request.dual_device_weekend_start_dates if item.strip()), ) try: schedule_error = scheduler.compile_schedule() entries = scheduler.get_schedule() except (RuntimeError, ValueError) as exc: raise HTTPException(status_code=409, detail=str(exc)) from exc completion_date = max((e.scheduled_date for e in entries), default=None) if schedule_error: raise HTTPException(status_code=409, detail=schedule_error) version = db.create_schedule_version( [ ( e.test_id, e.device, e.scheduled_date, e.shift_index, e.sequence_in_shift, status_lookup.get((e.test_id, e.device), "pending"), ) for e in entries ], DB_PATH, ) print(f"Schedule version {version} created with {len(entries)} entries, completion date: {completion_date}") return { "schedule_version": version, "scheduled_tests": len(entries), "completion_date": completion_date, } @app.get("/api/tests/rerun") def get_rerun_tests() -> dict[str, Any]: db.mark_overdue_as_rerun(DB_PATH) tests = db.get_rerun_tests(DB_PATH) total_minutes = sum(t["estimated_minutes"] for t in tests) return {"tests": tests, "total_estimated_minutes": total_minutes} @app.post("/api/holidays") def save_holidays(request: SaveHolidaysRequest) -> dict[str, Any]: dates = [d.strip() for d in request.dates if d.strip()] db.upsert_holidays(dates, DB_PATH) return {"status": "saved", "count": len(dates)} @app.get("/api/holidays") def get_holidays() -> dict[str, Any]: return {"dates": sorted(db.list_holidays(DB_PATH))} @app.get("/api/schedule/versions") def get_schedule_versions() -> dict[str, Any]: return {"versions": db.get_schedule_versions(DB_PATH)} @app.get("/api/schedule/week") def get_schedule_week(start: str | None = None, version: int | None = None) -> dict[str, Any]: week_start = start or date.today().isoformat() try: datetime.strptime(week_start, "%Y-%m-%d") except ValueError as exc: raise HTTPException(status_code=400, detail="start must be YYYY-MM-DD") from exc if version is not None and version <= 0: raise HTTPException(status_code=400, detail="version must be a positive integer") week_start_date = datetime.strptime(week_start, "%Y-%m-%d").date() selected_version = db.resolve_schedule_version(version, DB_PATH) if version is not None and selected_version is None: raise HTTPException(status_code=404, detail=f"Schedule version {version} was not found") rows = db.get_schedule_week(week_start, selected_version, DB_PATH) completed_rows = [row for row in rows if str(row.status).strip().lower() == "completed"] rerun_rows = [row for row in rows if str(row.status).strip().lower() == "rerun"] print( f"[api] /schedule/week: completed={len(completed_rows)} rerun_required={len(rerun_rows)}" ) all_rows = db.get_schedule_rows(selected_version, DB_PATH) completion_date = max((row.scheduled_date for row in all_rows), default=None) start_shift = min( ((row.scheduled_date, row.shift_index) for row in all_rows), default=(None, None), ) schedule_start_date, schedule_start_shift_index = start_shift holiday_dates = db.list_holidays(DB_PATH) return { "start_date": week_start, "schedule_version": selected_version, "schedule_start_date": schedule_start_date, "schedule_start_shift_index": schedule_start_shift_index, "items": [_serialize_schedule_row(row) for row in rows], "total_scheduled_tests": len(all_rows), "completion_date": completion_date, "windows": _build_schedule_windows(week_start_date, all_rows, holiday_dates), } if __name__ == "__main__": import uvicorn uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)