This commit is contained in:
@@ -0,0 +1,226 @@
|
||||
"""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 "")
|
||||
Reference in New Issue
Block a user