Files
scheduler/backend/parser.py
T
2026-07-13 16:17:36 -04:00

638 lines
18 KiB
Python

import csv
import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any
try:
import smbclient # type: ignore[import-not-found]
except ModuleNotFoundError:
smbclient = None
from db import DEVICE_DUT, DEVICE_REF, TestRecord
P2P_COE_REQUIRED_COLUMNS = [
"Priority",
"Index",
"Interferer",
"COE Pair",
"Rotation",
"TC ID",
"Victim Band",
"6GHz Power Mode",
"5G_Test Point",
"5G_Channel",
"5G_Bandwidth",
"5G_RSSI",
"5G_Direction",
"5G_STA",
"6G_Test Point",
"6G_Channel",
"6G_Bandwidth",
"6G_RSSI",
"6G_Direction",
"6G_STA",
"2G_Test Point",
"2G_Channel",
"2G_Bandwidth",
"2G_RSSI",
"2G_Direction",
"2G_STA",
]
P3P_REQUIRED_COLUMNS = [
"Priority",
"Index",
"Throttled",
"Rotation",
"TC ID",
"Band",
"6GHz Power Mode",
"STATION 1_Test Point",
"STATION 1_Channel",
"STATION 1_Bandwidth",
"STATION 1_RSSI",
"STATION 1_Direction",
"STATION 1_Rate",
"STATION 1_STA",
"STATION 2_Test Point",
"STATION 2_Channel",
"STATION 2_Bandwidth",
"STATION 2_RSSI",
"STATION 2_Direction",
"STATION 2_Rate",
"STATION 2_STA",
"STATION 3_Test Point",
"STATION 3_Channel",
"STATION 3_Bandwidth",
"STATION 3_RSSI",
"STATION 3_Direction",
"STATION 3_Rate",
"STATION 3_STA",
]
RUNTIME_DEFAULTS = {
"P2P": 80,
"COE": 115,
"P3P": 105,
}
@dataclass(frozen=True)
class ParseResult:
tests: list[TestRecord]
warnings: list[str]
class CsvValidationError(ValueError):
pass
_SMB_SESSIONS: set[str] = set()
def _normalize_input_path(path_value: str | Path) -> str:
path = str(path_value).strip()
if not path:
return path
# Accept //server/share style and normalize to UNC for smbclient.
if path.startswith("//"):
return "\\\\" + path.lstrip("/").replace("/", "\\")
# Accept /<ipv4>/<share>/... and normalize to UNC for Linux-hosted inputs.
if re.match(r"^/\d{1,3}(?:\.\d{1,3}){3}/[^/]+", path):
return "\\\\" + path.lstrip("/").replace("/", "\\")
return path
def _is_unc_path(path: str) -> bool:
return path.startswith("\\\\")
def _extract_unc_server(path: str) -> str | None:
if not _is_unc_path(path):
return None
rest = path[2:]
return rest.split("\\", 1)[0] if rest else None
def _normalize_smb_credentials(smb_credentials: dict[str, Any] | None) -> tuple[str, str, str]:
username = ""
password = ""
domain = ""
if smb_credentials:
username = str(smb_credentials.get("username", "")).strip()
password = str(smb_credentials.get("password", ""))
domain = str(smb_credentials.get("domain", "")).strip()
if not username:
username = os.getenv("SMB_USERNAME", "").strip()
if not password:
password = os.getenv("SMB_PASSWORD", "")
if not domain:
domain = os.getenv("SMB_DOMAIN", "").strip()
if username and domain and "\\" not in username and "@" not in username:
username = f"{domain}\\{username}"
return username, password, domain
def _register_smb_session_if_needed(path: str, smb_credentials: dict[str, Any] | None) -> None:
if not _is_unc_path(path):
return
if smbclient is None:
raise ModuleNotFoundError("smbclient is required to read CSV files from UNC paths")
server = _extract_unc_server(path)
if not server or server in _SMB_SESSIONS:
return
username, password, _ = _normalize_smb_credentials(smb_credentials)
if username:
smbclient.register_session(server, username=username, password=password)
else:
smbclient.register_session(server)
_SMB_SESSIONS.add(server)
def _path_exists(path: str, smb_credentials: dict[str, Any] | None) -> bool:
if _is_unc_path(path):
_register_smb_session_if_needed(path, smb_credentials)
try:
smbclient.stat(path)
return True
except OSError:
return False
return Path(path).exists()
def _is_dir(path: str, smb_credentials: dict[str, Any] | None) -> bool:
if _is_unc_path(path):
_register_smb_session_if_needed(path, smb_credentials)
try:
entry_iter = smbclient.scandir(path)
for _ in entry_iter:
break
return True
except OSError:
return False
return Path(path).is_dir()
def _path_name(path: str) -> str:
trimmed = path.rstrip("\\/")
if not trimmed:
return path
parts = re.split(r"[\\/]", trimmed)
return parts[-1] if parts else trimmed
def _csv_paths_from_dir(path: str, smb_credentials: dict[str, Any] | None) -> list[str]:
if _is_unc_path(path):
_register_smb_session_if_needed(path, smb_credentials)
entries = []
for entry in smbclient.scandir(path):
name = getattr(entry, "name", "")
if name and name.lower().endswith(".csv") and entry.is_file():
entries.append(path.rstrip("\\/") + "\\" + name)
return sorted(entries)
return [str(p) for p in sorted(Path(path).glob("*.csv"))]
def _columns_missing(normalized_fieldnames: set[str], required_columns: list[str]) -> list[str]:
return [
column
for column in required_columns
if _normalize_column_name(column) not in normalized_fieldnames
]
def _detect_csv_format(normalized_fieldnames: set[str]) -> str:
p2p_coe_missing = _columns_missing(normalized_fieldnames, P2P_COE_REQUIRED_COLUMNS)
if not p2p_coe_missing:
return "p2p_coe"
p3p_missing = _columns_missing(normalized_fieldnames, P3P_REQUIRED_COLUMNS)
if not p3p_missing:
return "p3p"
raise CsvValidationError(
"CSV format not recognized. Missing columns for P2P/COE format: "
f"{', '.join(p2p_coe_missing)}; missing columns for P3P format: {', '.join(p3p_missing)}"
)
def parse_target_csv(
csv_path: str | Path | list[str | Path] | tuple[str | Path, ...],
smb_credentials: dict[str, Any] | None = None,
runtime_overrides: dict[str, Any] | None = None,
) -> ParseResult:
paths = _resolve_csv_paths(csv_path, smb_credentials=smb_credentials)
runtime_defaults_by_device = _resolve_runtime_defaults(runtime_overrides)
all_tests: list[TestRecord] = []
all_warnings: list[str] = []
seen_test_keys: set[tuple[str, str]] = set()
for path in paths:
parsed = _parse_single_csv(
path,
smb_credentials=smb_credentials,
runtime_defaults_by_device=runtime_defaults_by_device,
)
all_warnings.extend(parsed.warnings)
for record in parsed.tests:
key = (record.test_id, record.device)
if key in seen_test_keys:
all_warnings.append(
f"File {_path_name(path)}: duplicate test/device '{record.test_id}/{record.device}', record skipped."
)
continue
seen_test_keys.add(key)
all_tests.append(record)
return ParseResult(tests=all_tests, warnings=all_warnings)
def _coerce_runtime_defaults(raw_overrides: dict[str, Any] | None) -> dict[str, int]:
defaults = dict(RUNTIME_DEFAULTS)
if not raw_overrides:
return defaults
for test_type in ("P2P", "COE", "P3P"):
raw_value = raw_overrides.get(test_type)
if raw_value is None:
continue
if isinstance(raw_value, str):
raw_value = raw_value.strip()
if not raw_value:
continue
try:
minutes = int(raw_value)
except (TypeError, ValueError):
continue
if minutes > 0:
defaults[test_type] = minutes
return defaults
def _resolve_runtime_defaults(runtime_overrides: dict[str, Any] | None) -> dict[str, dict[str, int]]:
legacy_defaults = _coerce_runtime_defaults(runtime_overrides)
if not runtime_overrides:
return {
DEVICE_DUT: dict(legacy_defaults),
DEVICE_REF: dict(legacy_defaults),
}
resolved: dict[str, dict[str, int]] = {}
for device in (DEVICE_DUT, DEVICE_REF):
device_defaults = dict(legacy_defaults)
raw_device_overrides = runtime_overrides.get(device)
if isinstance(raw_device_overrides, dict):
device_defaults.update(_coerce_runtime_defaults(raw_device_overrides))
resolved[device] = device_defaults
return resolved
def _resolve_csv_paths(
csv_path: str | Path | list[str | Path] | tuple[str | Path, ...],
smb_credentials: dict[str, Any] | None,
) -> list[str]:
if isinstance(csv_path, (str, Path)):
paths = [_normalize_input_path(csv_path)]
else:
paths = [_normalize_input_path(item) for item in csv_path]
resolved_paths: list[str] = []
for path in paths:
if _is_dir(path, smb_credentials=smb_credentials):
resolved_paths.extend(_csv_paths_from_dir(path, smb_credentials=smb_credentials))
else:
resolved_paths.append(path)
if not resolved_paths:
raise FileNotFoundError("No CSV files found to parse.")
for path in resolved_paths:
if not _path_exists(path, smb_credentials=smb_credentials):
raise FileNotFoundError(f"CSV file not found: {path}")
return resolved_paths
def _open_csv_handle(path: str, smb_credentials: dict[str, Any] | None):
if _is_unc_path(path):
_register_smb_session_if_needed(path, smb_credentials)
return smbclient.open_file(path, mode="r", encoding="utf-8-sig", newline="")
return Path(path).open("r", encoding="utf-8-sig", newline="")
def _parse_single_csv(
path: str,
smb_credentials: dict[str, Any] | None,
runtime_defaults_by_device: dict[str, dict[str, int]],
) -> ParseResult:
with _open_csv_handle(path, smb_credentials=smb_credentials) as handle:
reader = csv.DictReader(handle)
if not reader.fieldnames:
raise CsvValidationError("CSV is missing a header row.")
normalized_fieldnames = {_normalize_column_name(name) for name in reader.fieldnames if name}
csv_format = _detect_csv_format(normalized_fieldnames)
tests: list[TestRecord] = []
warnings: list[str] = []
seen_test_ids: set[str] = set()
records_with_signature: list[tuple[TestRecord, tuple[str, ...] | None]] = []
for row_num, row in enumerate(reader, start=2):
test_id = (row.get("TC ID") or "").strip()
if not test_id:
warnings.append(f"File {_path_name(path)}, row {row_num}: missing TC ID, row skipped.")
continue
if test_id in seen_test_ids:
warnings.append(f"File {_path_name(path)}, row {row_num}: duplicate TC ID '{test_id}', row skipped.")
continue
seen_test_ids.add(test_id)
throttled = _normalize_yes_no(_row_get(row, "Throttled")) if csv_format == "p3p" else False
test_type = _infer_test_type(test_id)
rx_tx = _infer_rx_tx(test_id)
rotation = _empty_to_none(_row_get(row, "Rotation"))
has_coe_pair = _normalize_yes_no(_row_get(row, "COE Pair")) if csv_format == "p2p_coe" else False
config = _build_config(row, csv_format)
signature = _victim_band_signature(row) if csv_format == "p2p_coe" else None
victim_band_source = "Victim Band" if csv_format == "p2p_coe" else "Band"
victim_band = _normalize_victim_band(_row_get(row, victim_band_source))
# Default priority based on test type and COE pairing.
if test_type == "P2P":
priority = 2 if has_coe_pair else 3
elif test_type == "COE":
priority = 3
else:
priority = 4
record = TestRecord(
test_id=test_id,
device=DEVICE_DUT,
test_type=test_type,
rotation=rotation,
rx_tx=rx_tx,
has_coe_pair=has_coe_pair,
coe_pairing=[],
priority=priority,
victim_band=victim_band,
config=config,
throttled=throttled,
estimated_minutes=RUNTIME_DEFAULTS[test_type],
status="pending",
excluded=False,
raw_payload=row,
)
records_with_signature.append((record, signature))
# Build COE pairing based on victim band signature and RX/TX+band suffix.
# Example key suffixes: RXAX, TXAX, RXBE, TXBE.
coe_by_signature_and_suffix: dict[tuple, list[str]] = {}
for record, signature in records_with_signature:
if record.test_type == "COE" and signature is not None:
suffix = _extract_pairing_suffix(record.test_id)
if suffix:
key = (signature, suffix)
coe_by_signature_and_suffix.setdefault(key, []).append(record.test_id)
for record, signature in records_with_signature:
pairs = record.coe_pairing
if record.test_type == "P2P" and signature is not None:
suffix = _extract_pairing_suffix(record.test_id)
key = (signature, suffix) if suffix else None
pairs = sorted(coe_by_signature_and_suffix.get(key, [])) if key else []
for device in (DEVICE_DUT, DEVICE_REF):
device_runtime_defaults = runtime_defaults_by_device.get(device, RUNTIME_DEFAULTS)
tests.append(
TestRecord(
test_id=record.test_id,
device=device,
test_type=record.test_type,
rotation=record.rotation,
rx_tx=record.rx_tx,
has_coe_pair=bool(pairs),
coe_pairing=pairs,
priority=record.priority,
victim_band=record.victim_band,
config=record.config,
throttled=record.throttled,
estimated_minutes=device_runtime_defaults.get(record.test_type, RUNTIME_DEFAULTS[record.test_type]),
status=record.status,
excluded=record.excluded,
raw_payload=record.raw_payload,
)
)
return ParseResult(tests=tests, warnings=warnings)
def _infer_test_type(test_id: str) -> str:
token = test_id.upper()
if token.startswith("COE"):
return "COE"
if token.startswith("P3P"):
return "P3P"
return "P2P"
def _infer_rx_tx(test_id: str) -> str | None:
token = test_id.upper()
if "RX" in token:
return "RX"
if "TX" in token:
return "TX"
return None
def _extract_pairing_suffix(test_id: str) -> str | None:
"""Extract the RX/TX+band suffix used for pairing.
Format: [COE|P2P][RX|TX][Band][Number]
Example: P2PRXAC004 -> 'RXAC', COETXAX012 -> 'TXAX'
"""
token = test_id.upper()
# Remove COE/P2P prefix
if token.startswith("COE"):
token = token[3:]
elif token.startswith("P2P") or token.startswith("P3P"):
token = token[3:]
else:
return None
# Extract RX/TX + band (for example RXAX, TXBE).
if len(token) >= 4 and (token.startswith("RX") or token.startswith("TX")):
return token[:4]
return None
def _empty_to_none(value: str | None) -> str | None:
if value is None:
return None
cleaned = value.strip()
return cleaned if cleaned else None
def _victim_band_signature(row: dict[str, str]) -> tuple[str, ...] | None:
band = _normalize_victim_band(_row_get(row, "Victim Band"))
if band is None:
return None
# Source CSVs may use either "5G_Test Point" or "5G Test Point" style headers.
prefixes = (f"{band}_", f"{band} ")
def _pick_value(suffix: str) -> str:
for prefix in prefixes:
value = _row_get(row, f"{prefix}{suffix}")
if _normalize_value(value):
return _normalize_value(value)
return ""
test_point = _pick_value("Test Point")
channel = _pick_value("Channel")
rssi = _pick_value("RSSI")
bandwidth = _pick_value("Bandwidth")
# Direction is intentionally ignored for COE pairing matching.
if not all([test_point, channel, rssi, bandwidth]):
return None
return (band, test_point, channel, rssi, bandwidth)
def _normalize_victim_band(value: str | None) -> str | None:
if value is None:
return None
normalized = value.strip().upper().replace("GHZ", "G")
if normalized.startswith("5"):
return "5G"
if normalized.startswith("6"):
return "6G"
if normalized.startswith("2"):
return "2G"
return None
def _normalize_value(value: str | None) -> str:
if value is None:
return ""
return " ".join(value.strip().upper().split())
def _normalize_column_name(value: str | None) -> str:
if value is None:
return ""
# Accept common header variants, for example: "TC ID", "TC_ID", or "TC\u00A0ID".
return re.sub(r"[^A-Z0-9]+", "", value.upper())
def _normalize_yes_no(value: str | None) -> bool:
return _normalize_value(value) == "YES"
def _build_legacy_config(row: dict[str, str]) -> dict[str, dict[str, str | None]]:
return {
"5G": {
"test_point": _empty_to_none(_row_get(row, "5G_Test Point")),
"channel": _empty_to_none(_row_get(row, "5G_Channel")),
"bandwidth": _empty_to_none(_row_get(row, "5G_Bandwidth")),
"rssi": _empty_to_none(_row_get(row, "5G_RSSI")),
"direction": _empty_to_none(_row_get(row, "5G_Direction")),
"sta": _empty_to_none(_row_get(row, "5G_STA")),
},
"6G": {
"power_mode": _empty_to_none(_row_get(row, "6GHz Power Mode")),
"test_point": _empty_to_none(_row_get(row, "6G_Test Point")),
"channel": _empty_to_none(_row_get(row, "6G_Channel")),
"bandwidth": _empty_to_none(_row_get(row, "6G_Bandwidth")),
"rssi": _empty_to_none(_row_get(row, "6G_RSSI")),
"direction": _empty_to_none(_row_get(row, "6G_Direction")),
"sta": _empty_to_none(_row_get(row, "6G_STA")),
},
"2G": {
"test_point": _empty_to_none(_row_get(row, "2G_Test Point")),
"channel": _empty_to_none(_row_get(row, "2G_Channel")),
"bandwidth": _empty_to_none(_row_get(row, "2G_Bandwidth")),
"rssi": _empty_to_none(_row_get(row, "2G_RSSI")),
"direction": _empty_to_none(_row_get(row, "2G_Direction")),
"sta": _empty_to_none(_row_get(row, "2G_STA")),
},
}
def _build_p3p_config(row: dict[str, str]) -> dict[str, dict[str, str | None]]:
return {
"STATION1": {
"test_point": _empty_to_none(_row_get(row, "STATION 1_Test Point")),
"channel": _empty_to_none(_row_get(row, "STATION 1_Channel")),
"bandwidth": _empty_to_none(_row_get(row, "STATION 1_Bandwidth")),
"rssi": _empty_to_none(_row_get(row, "STATION 1_RSSI")),
"direction": _empty_to_none(_row_get(row, "STATION 1_Direction")),
"rate": _empty_to_none(_row_get(row, "STATION 1_Rate")),
"sta": _empty_to_none(_row_get(row, "STATION 1_STA")),
},
"STATION2": {
"test_point": _empty_to_none(_row_get(row, "STATION 2_Test Point")),
"channel": _empty_to_none(_row_get(row, "STATION 2_Channel")),
"bandwidth": _empty_to_none(_row_get(row, "STATION 2_Bandwidth")),
"rssi": _empty_to_none(_row_get(row, "STATION 2_RSSI")),
"direction": _empty_to_none(_row_get(row, "STATION 2_Direction")),
"rate": _empty_to_none(_row_get(row, "STATION 2_Rate")),
"sta": _empty_to_none(_row_get(row, "STATION 2_STA")),
},
"STATION3": {
"test_point": _empty_to_none(_row_get(row, "STATION 3_Test Point")),
"channel": _empty_to_none(_row_get(row, "STATION 3_Channel")),
"bandwidth": _empty_to_none(_row_get(row, "STATION 3_Bandwidth")),
"rssi": _empty_to_none(_row_get(row, "STATION 3_RSSI")),
"direction": _empty_to_none(_row_get(row, "STATION 3_Direction")),
"rate": _empty_to_none(_row_get(row, "STATION 3_Rate")),
"sta": _empty_to_none(_row_get(row, "STATION 3_STA")),
},
}
def _build_config(row: dict[str, str], csv_format: str) -> dict[str, dict[str, str | None]]:
if csv_format == "p3p":
return _build_p3p_config(row)
return _build_legacy_config(row)
def _row_get(row: dict[str, str], key: str) -> str | None:
if key in row:
return row.get(key)
normalized_key = _normalize_column_name(key)
for existing_key, value in row.items():
if existing_key and _normalize_column_name(existing_key) == normalized_key:
return value
if normalized_key == _normalize_column_name("COE PAIR"):
for alias in ("COE PAIRING", "COE_PAIRING"):
for existing_key, value in row.items():
if existing_key and _normalize_column_name(existing_key) == _normalize_column_name(alias):
return value
return None