"""Tests for main module.""" from __future__ import annotations from pathlib import Path from unittest.mock import MagicMock, patch from main import ( RuleConfig, build_zone_map, load_config, process_rule, send_ntfy_notification, ) from unifi_firewall import FirewallZone class TestBuildZoneMap: def test_builds_map_from_zones(self) -> None: zones: list[FirewallZone] = [ {"id": "zone-wan", "name": "WAN", "networkIds": [], "metadata": {}}, {"id": "zone-lan", "name": "LAN", "networkIds": [], "metadata": {}}, ] zone_map = build_zone_map(zones) assert zone_map == {"WAN": "zone-wan", "LAN": "zone-lan"} def test_empty_zones(self) -> None: zone_map = build_zone_map([]) assert zone_map == {} class TestLoadConfig: def test_loads_valid_config(self, tmp_path: Path) -> None: config_file = tmp_path / "rules.yaml" config_file.write_text( """ rules: - name: "Test Rule" source_zone: "WAN" dest_zone: "LAN" ip_version: "IPV4" action: "ALLOW" """ ) rules = load_config(str(config_file)) assert len(rules) == 1 assert rules[0].get("name") == "Test Rule" def test_exits_on_missing_file(self) -> None: with patch("sys.exit") as mock_exit: load_config("/nonexistent/rules.yaml") mock_exit.assert_called_once_with(1) def test_exits_on_invalid_yaml(self, tmp_path: Path) -> None: config_file = tmp_path / "rules.yaml" config_file.write_text("invalid: yaml: content: [") with patch("sys.exit") as mock_exit: load_config(str(config_file)) mock_exit.assert_called_once_with(1) class TestSendNtfyNotification: def test_sends_notification(self) -> None: mock_response = MagicMock() mock_response.raise_for_status = MagicMock() with ( patch("main.NTFY_URL", "https://ntfy.sh"), patch("main.NTFY_TOPIC", "test"), patch("main.NTFY_API_KEY", "test-key"), patch("main.requests.post", return_value=mock_response) as mock_post, ): send_ntfy_notification("Title", "Message", priority=4) mock_post.assert_called_once() call_kwargs = mock_post.call_args assert call_kwargs[0][0] == "https://ntfy.sh/test" assert call_kwargs[1]["data"] == b"Message" assert call_kwargs[1]["headers"]["Title"] == "Title" assert call_kwargs[1]["headers"]["Priority"] == "4" def test_skips_when_no_url(self) -> None: with ( patch("main.NTFY_URL", ""), patch("main.requests.post") as mock_post, ): send_ntfy_notification("Title", "Message") mock_post.assert_not_called() def test_handles_request_error(self) -> None: import requests as req with ( patch("main.NTFY_URL", "https://ntfy.sh"), patch("main.NTFY_TOPIC", "test"), patch("main.requests.post", side_effect=req.RequestException("timeout")), ): # Should not raise, only log warning send_ntfy_notification("Title", "Message") class TestProcessRule: def test_creates_new_rule(self) -> None: mock_session = MagicMock() mock_session.get.return_value = MagicMock(json=MagicMock(return_value={"data": []}), raise_for_status=MagicMock()) mock_session.post.return_value = MagicMock( json=MagicMock(return_value={"id": "new-id", "name": "Test Rule"}), raise_for_status=MagicMock(), ) rule: RuleConfig = { "name": "Test Rule", "source_zone": "WAN", "dest_zone": "LAN", "ip_version": "IPV4", "action": "ALLOW", } zone_map = {"WAN": "wan-id", "LAN": "lan-id"} change = process_rule(mock_session, "https://unifi.local", "site-1", rule, zone_map, "1.2.3.4") assert change["action"] == "created" assert change["rule_name"] == "Test Rule" assert change["error"] is None def test_skips_rule_with_matching_ip(self) -> None: mock_session = MagicMock() existing_policy = { "id": "pol-1", "name": "Test Rule", "source": { "zoneId": "wan-id", "trafficFilter": { "type": "IP_ADDRESS", "ipAddressFilter": { "type": "IP_ADDRESSES", "items": [{"type": "IP_ADDRESS", "value": "1.2.3.4"}], }, }, }, } mock_session.get.return_value = MagicMock(json=MagicMock(return_value={"data": [existing_policy]}), raise_for_status=MagicMock()) rule: RuleConfig = { "name": "Test Rule", "source_zone": "WAN", "dest_zone": "LAN", "ip_version": "IPV4", } zone_map = {"WAN": "wan-id", "LAN": "lan-id"} change = process_rule(mock_session, "https://unifi.local", "site-1", rule, zone_map, "1.2.3.4") assert change["action"] == "skipped" mock_session.put.assert_not_called() mock_session.post.assert_not_called() def test_updates_rule_with_different_ip(self) -> None: mock_session = MagicMock() existing_policy = { "id": "pol-1", "name": "Test Rule", "source": { "zoneId": "wan-id", "trafficFilter": { "type": "IP_ADDRESS", "ipAddressFilter": { "type": "IP_ADDRESSES", "items": [{"type": "IP_ADDRESS", "value": "5.6.7.8"}], }, }, }, } mock_session.get.return_value = MagicMock(json=MagicMock(return_value={"data": [existing_policy]}), raise_for_status=MagicMock()) mock_session.put.return_value = MagicMock(json=MagicMock(return_value={"id": "pol-1"}), raise_for_status=MagicMock()) rule: RuleConfig = { "name": "Test Rule", "source_zone": "WAN", "dest_zone": "LAN", "ip_version": "IPV4", } zone_map = {"WAN": "wan-id", "LAN": "lan-id"} change = process_rule(mock_session, "https://unifi.local", "site-1", rule, zone_map, "1.2.3.4") assert change["action"] == "updated" mock_session.put.assert_called_once() def test_fails_on_missing_source_zone(self) -> None: mock_session = MagicMock() rule: RuleConfig = { "name": "Test Rule", "source_zone": "UNKNOWN", "dest_zone": "LAN", "ip_version": "IPV4", } zone_map = {"LAN": "lan-id"} change = process_rule(mock_session, "https://unifi.local", "site-1", rule, zone_map, "1.2.3.4") assert change["action"] == "failed" assert "Source zone" in (change["error"] or "") def test_fails_on_api_error(self) -> None: mock_session = MagicMock() mock_session.get.side_effect = Exception("Connection error") rule: RuleConfig = { "name": "Test Rule", "source_zone": "WAN", "dest_zone": "LAN", "ip_version": "IPV4", } zone_map = {"WAN": "wan-id", "LAN": "lan-id"} change = process_rule(mock_session, "https://unifi.local", "site-1", rule, zone_map, "1.2.3.4") assert change["action"] == "failed" assert "Failed to list policies" in (change["error"] or "")