scheduler/backend/test_config.py

from __future__ import annotations

import json
from typing import Any

TEST_CONFIG_DEFINITION = {
	"TC1": {"T1D": "STA5", "T1F": "STA56", "T1I": "STA58", "T1L": "STA64", "T2A":"STA63", "T2O":"STA65", "T2J": "STA59", "T2E": "STA6", "T3E": "STA4"},
	"TC2": {"T1D": "STA64", "T1F": "STA4", "T1I": "STA5", "T1L": "STA58", "T2A":"STA56", "T2O":"STA59", "T2J": "STA6", "T2E": "STA65", "T3E": "STA63"},
	"TC3": {"T1D": "STA58", "T1F": "STA63", "T1I": "STA64", "T2J": "STA65", "T2E": "STA59", "T3E": "STA56"},
	"TC4": {"T1B": "STA56", "T1C": "STA4", "T1D": "STA63"},
	"TC5": {"T1B": "STA4", "T1C": "STA63", "T1D": "STA56"},
	"TC6": {"T1B": "STA63", "T1C": "STA56", "T1D": "STA4"},
	"TC7": {"T1F": "STA4", "T2A": "STA63", "T3E": "STA56"},
	"TC8": {"T1F": "STA63", "T2A": "STA56", "T3E": "STA4"},
	"TC9": {"T1F": "STA63", "T2A": "STA64", "T3E": "STA4"},
	"TC10": {"T1F": "STA4", "T3E": "STA56", "T1K2A": "STA63"},
	"TC11": {"T1F": "STA56", "T3E": "STA63", "T1K2A": "STA4"},
	"TC12": {"T1F": "STA56", "T3E": "STA4", "T1K2A": "STA63"},
}



def _norm(value: Any) -> str:
	if value is None:
		return ""
	return " ".join(str(value).strip().upper().split())


def _get_test_config(test: Any) -> dict[str, dict[str, str | None]]:
	if isinstance(test, dict):
		config = test.get("config") or {}
	else:
		config = getattr(test, "config", None) or {}
	return config if isinstance(config, dict) else {}


def _testpoint_to_station_sets(config: dict[str, dict[str, str | None]]) -> dict[str, set[str]]:
	result: dict[str, set[str]] = {}

	def _add_entry(entry: dict[str, str | None]) -> None:
		testpoint = _norm(entry.get("test_point") or entry.get("Test Point"))
		sta_raw = _norm(entry.get("sta") or entry.get("STA"))
		if not testpoint or not sta_raw:
			return

		stations = {_norm(sta) for sta in sta_raw.split(",") if _norm(sta)}
		if not stations:
			return

		result.setdefault(testpoint, set()).update(stations)

	for entry in config.values():
		if isinstance(entry, dict):
			_add_entry(entry)

	return result


def build_station_testpoint_map(config: dict[str, dict[str, str | None]]) -> dict[str, str]:
	"""Build station->testpoint map used by legacy graph serialization."""
	station_to_testpoint: dict[str, str] = {}
	for testpoint, stations in _testpoint_to_station_sets(config).items():
		for station in stations:
			station_to_testpoint[station] = testpoint
	return station_to_testpoint


def resolve_test_config_keys(config: dict[str, dict[str, str | None]]) -> list[str]:
	"""Return all TC keys whose definitions contain all testpoint->station mappings in config."""
	testpoint_stations = _testpoint_to_station_sets(config)
	if not testpoint_stations:
		return []

	matches: list[str] = []
	for tc_key, tc_mapping in TEST_CONFIG_DEFINITION.items():
		is_match = True
		for testpoint, stations in testpoint_stations.items():
			# TC definitions map one testpoint to exactly one station.
			if len(stations) != 1:
				is_match = False
				break
			expected_station = tc_mapping.get(testpoint)
			if expected_station is None or expected_station not in stations:
				is_match = False
				break
		if is_match:
			matches.append(tc_key)

	return matches


def build_config_rows(tests: list[dict[str, Any]]) -> list[dict[str, Any]]:
	by_sta: dict[str, set[str]] = {}

	for test in tests:
		config = test.get("config") or {}
		if not isinstance(config, dict):
			continue
		for entry in config.values():
			if not isinstance(entry, dict):
				continue
			sta_raw = entry.get("sta") or entry.get("STA")
			test_point = entry.get("test_point") or entry.get("Test Point")
			if not sta_raw or not test_point:
				continue

			for sta_part in str(sta_raw).split(","):
				sta = sta_part.strip().upper()
				if not sta:
					continue
				by_sta.setdefault(sta, set()).add(str(test_point))

	return [
		{
			"sta": sta,
			"testPoints": sorted(points),
		}
		for sta, points in sorted(by_sta.items())
	]


def build_bundle_test_configs(tests: list[Any]) -> list[str]:
	merged_by_testpoint: dict[str, set[str]] = {}
	for test in tests:
		config = _get_test_config(test)
		for entry in config.values():
			if not isinstance(entry, dict):
				continue
			testpoint_raw = entry.get("test_point") or entry.get("Test Point")
			sta_raw = entry.get("sta") or entry.get("STA")
			if not testpoint_raw or not sta_raw:
				continue

			testpoint = " ".join(str(testpoint_raw).strip().upper().split())
			if not testpoint:
				continue

			stations = {
				" ".join(str(sta).strip().upper().split())
				for sta in str(sta_raw).split(",")
				if " ".join(str(sta).strip().upper().split())
			}
			if not stations:
				continue

			merged_by_testpoint.setdefault(testpoint, set()).update(stations)

	merged_config: dict[str, dict[str, str | None]] = {}
	for idx, (testpoint, stations) in enumerate(sorted(merged_by_testpoint.items()), start=1):
		merged_config[f"Station {idx}"] = {
			"test_point": testpoint,
			"sta": ",".join(sorted(stations)),
		}

	return resolve_test_config_keys(merged_config)


def serialize_station_testpoint_map(config: dict[str, dict[str, str | None]]) -> str:
	return json.dumps(build_station_testpoint_map(config))