import unittest
from pathlib import Path

import yaml


CONFIG_DIR = Path(__file__).resolve().parents[1] / "config"


def _duration_by_node(template_path: Path) -> dict[str, object]:
    workflow = yaml.safe_load(template_path.read_text())
    durations = {}

    for node in workflow.get("nodes", []):
        name = node.get("data", {}).get("name")
        node_type = node.get("type")
        spec = node.get("data", {}).get("spec", {})

        if node_type == "sleep":
            durations[name] = spec.get("duration")
        elif node_type == "ti_workload_call" and spec.get("workload"):
            workload = yaml.safe_load(spec["workload"])
            duration = workload.get("params", {}).get("duration")
            if duration is not None:
                durations[name] = duration

    return durations


class WorkflowTemplateParityTests(unittest.TestCase):
    def test_premium_sysbench_durations_match_dedicated(self):
        dedicated = _duration_by_node(CONFIG_DIR / "perf_sysbench.yaml")
        premium = _duration_by_node(CONFIG_DIR / "perf_x_sysbench.yaml")

        self.assertEqual(premium, dedicated)

    def test_premium_tpcc5k_durations_match_dedicated(self):
        dedicated = _duration_by_node(CONFIG_DIR / "perf_tpcc5k.yaml")
        premium = _duration_by_node(CONFIG_DIR / "perf_x_tpcc5k.yaml")

        self.assertEqual(premium, dedicated)


if __name__ == "__main__":
    unittest.main()