tests/test_site_counts.py

#!/usr/bin/env python3
"""
Unit tests for site count parsing and topology counting helpers.
"""

import unittest
from agent.pipeline_engine import OvergrowthPipeline


class SiteCountTests(unittest.TestCase):
    def setUp(self):
        # Keep NetBox disabled for fast, offline unit tests.
        self.pipeline = OvergrowthPipeline(use_netbox=False)

    def test_parse_site_count_variants(self):
        self.assertEqual(
            self.pipeline._parse_site_count("Design a 4-site retail network with WiFi"),
            4,
        )
        self.assertEqual(
            self.pipeline._parse_site_count("We need a seven site deployment"), 7
        )
        self.assertEqual(
            self.pipeline._parse_site_count("Build 9 stores across the region"), 9
        )

    def test_site_names_alignment(self):
        self.assertEqual(
            self.pipeline._site_names(3),
            ["☕-Store-1", "☕-Store-2", "☕-Store-3"],
        )

    def test_count_branch_sites_exact(self):
        nodes = [
            {"name": "☕-Store-1", "node_type": "cloud"},
            {"name": "☕-Store-2", "node_type": "cloud"},
            {"name": "SW-Store-1", "node_type": "qemu"},
            {"name": "SW-Store-2", "node_type": "qemu"},
            {"name": "☁️-Internet-MPLS", "node_type": "cloud"},
            {"name": "-HQ-Building", "node_type": "cloud"},
        ]
        summary = self.pipeline._count_branch_sites(nodes, expected_site_count=2)
        self.assertEqual(summary["count"], 2)
        self.assertEqual(summary["missing"], [])

    def test_count_branch_sites_missing(self):
        nodes = [
            {"name": "☕-Store-1", "node_type": "cloud"},
            {"name": "☕-Store-2", "node_type": "cloud"},
            {"name": "☕-Store-3", "node_type": "cloud"},
            {"name": "SW-Store-1", "node_type": "qemu"},
            {"name": "SW-Store-2", "node_type": "qemu"},
            {"name": "SW-Store-3", "node_type": "qemu"},
        ]
        summary = self.pipeline._count_branch_sites(nodes, expected_site_count=4)
        self.assertEqual(summary["count"], 3)
        self.assertEqual(summary["missing"], ["☕-Store-4"])


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