init
Build and Push Container / build-and-push (push) Successful in 19s

This commit is contained in:
2026-08-01 19:07:05 -04:00
commit e360546de9
17 changed files with 15731 additions and 0 deletions
+47
View File
@@ -0,0 +1,47 @@
"""Tests for ip_lookup module."""
from unittest.mock import MagicMock, patch
from ip_lookup import get_ipv4, get_ipv6
class TestGetIpv4:
def test_returns_stripped_ip(self) -> None:
mock_result = MagicMock()
mock_result.stdout = " 1.2.3.4 \n"
with patch("ip_lookup.subprocess.run", return_value=mock_result) as mock_run:
ip = get_ipv4()
assert ip == "1.2.3.4"
mock_run.assert_called_once_with(["curl", "-4", "ifconfig.me"], capture_output=True, text=True, check=False)
def test_empty_response(self) -> None:
mock_result = MagicMock()
mock_result.stdout = ""
with patch("ip_lookup.subprocess.run", return_value=mock_result):
ip = get_ipv4()
assert ip == ""
class TestGetIpv6:
def test_returns_stripped_ip(self) -> None:
mock_result = MagicMock()
mock_result.stdout = " 2001:db8::1 \n"
with patch("ip_lookup.subprocess.run", return_value=mock_result) as mock_run:
ip = get_ipv6()
assert ip == "2001:db8::1"
mock_run.assert_called_once_with(["curl", "-6", "ifconfig.me"], capture_output=True, text=True, check=False)
def test_empty_response(self) -> None:
mock_result = MagicMock()
mock_result.stdout = ""
with patch("ip_lookup.subprocess.run", return_value=mock_result):
ip = get_ipv6()
assert ip == ""
+226
View File
@@ -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 "")
+299
View File
@@ -0,0 +1,299 @@
"""Tests for unifi_firewall module."""
from __future__ import annotations
from typing import Any
from unittest.mock import MagicMock
from unifi_firewall import (
FirewallPolicyCreatePayload,
build_policy_payload,
create_policy,
get_policy,
get_session,
get_source_ip_from_policy,
list_policies,
list_zones,
update_policy,
)
class MockHttpError(Exception):
"""Mock HTTP error for testing."""
class MockResponse:
def __init__(self, json_data: Any, status_code: int = 200) -> None:
self.json_data = json_data
self.status_code = status_code
def json(self) -> Any:
return self.json_data
def raise_for_status(self) -> None:
if self.status_code >= 400:
raise MockHttpError(f"HTTP {self.status_code}")
class TestGetSession:
def test_creates_session_with_headers(self) -> None:
session = get_session("https://unifi.local", "test-token", True)
assert session.headers["X-API-Key"] == "test-token"
assert session.headers["Content-Type"] == "application/json"
assert session.verify is True
def test_verify_ssl_false(self) -> None:
session = get_session("https://unifi.local", "test-token", False)
assert session.verify is False
class TestListZones:
def test_returns_zones_from_data_key(self) -> None:
mock_session = MagicMock()
mock_session.get.return_value = MockResponse({"data": [{"id": "zone-1", "name": "WAN", "networkIds": [], "metadata": {}}]})
zones = list_zones(mock_session, "https://unifi.local", "site-1")
assert len(zones) == 1
assert zones[0]["name"] == "WAN"
mock_session.get.assert_called_once_with("https://unifi.local/proxy/network/integration/v1/sites/site-1/firewall/zones")
def test_returns_zones_as_list(self) -> None:
mock_session = MagicMock()
mock_session.get.return_value = MockResponse([{"id": "zone-1", "name": "LAN", "networkIds": [], "metadata": {}}])
zones = list_zones(mock_session, "https://unifi.local", "site-1")
assert len(zones) == 1
assert zones[0]["name"] == "LAN"
class TestListPolicies:
def test_returns_policies_without_filter(self) -> None:
mock_session = MagicMock()
mock_session.get.return_value = MockResponse({"data": []})
policies = list_policies(mock_session, "https://unifi.local", "site-1")
assert policies == []
mock_session.get.assert_called_once_with("https://unifi.local/proxy/network/integration/v1/sites/site-1/firewall/policies", params={})
def test_returns_policies_with_name_filter(self) -> None:
mock_session = MagicMock()
mock_session.get.return_value = MockResponse({"data": []})
policies = list_policies(mock_session, "https://unifi.local", "site-1", name_filter="Test Rule")
assert policies == []
call_args = mock_session.get.call_args
assert call_args[1]["params"]["filter"] == "name.like('Test Rule')"
class TestGetPolicy:
def test_returns_single_policy(self) -> None:
mock_session = MagicMock()
policy_data = {"id": "pol-1", "name": "Test", "source": {"zoneId": "z1", "trafficFilter": None}, "destination": {"zoneId": "z2", "trafficFilter": None}, "action": {"type": "ALLOW", "allowReturnTraffic": True}, "ipProtocolScope": {"ipVersion": "IPV4"}, "enabled": True, "loggingEnabled": False, "index": 0, "metadata": {}}
mock_session.get.return_value = MockResponse(policy_data)
policy = get_policy(mock_session, "https://unifi.local", "site-1", "pol-1")
assert policy["id"] == "pol-1"
mock_session.get.assert_called_once_with("https://unifi.local/proxy/network/integration/v1/sites/site-1/firewall/policies/pol-1")
class TestCreatePolicy:
def test_creates_policy(self) -> None:
mock_session = MagicMock()
result = {"id": "new-pol", "name": "New Rule", "source": {"zoneId": "z1", "trafficFilter": None}, "destination": {"zoneId": "z2", "trafficFilter": None}, "action": {"type": "ALLOW", "allowReturnTraffic": True}, "ipProtocolScope": {"ipVersion": "IPV4"}, "enabled": True, "loggingEnabled": False, "index": 0, "metadata": {}}
mock_session.post.return_value = MockResponse(result, status_code=201)
payload: FirewallPolicyCreatePayload = build_policy_payload(
name="New Rule",
source_ip="1.2.3.4",
ip_version="IPV4",
src_zone_id="z1",
dst_zone_id="z2",
action_type="ALLOW",
)
policy = create_policy(mock_session, "https://unifi.local", "site-1", payload)
assert policy["id"] == "new-pol"
mock_session.post.assert_called_once()
class TestUpdatePolicy:
def test_updates_policy(self) -> None:
mock_session = MagicMock()
result = {"id": "pol-1", "name": "Updated Rule", "source": {"zoneId": "z1", "trafficFilter": None}, "destination": {"zoneId": "z2", "trafficFilter": None}, "action": {"type": "ALLOW", "allowReturnTraffic": True}, "ipProtocolScope": {"ipVersion": "IPV4"}, "enabled": True, "loggingEnabled": False, "index": 0, "metadata": {}}
mock_session.put.return_value = MockResponse(result)
payload: FirewallPolicyCreatePayload = build_policy_payload(
name="Updated Rule",
source_ip="5.6.7.8",
ip_version="IPV4",
src_zone_id="z1",
dst_zone_id="z2",
action_type="ALLOW",
)
policy = update_policy(mock_session, "https://unifi.local", "site-1", "pol-1", payload)
assert policy["id"] == "pol-1"
mock_session.put.assert_called_once()
class TestBuildPolicyPayload:
def test_basic_allow_ipv4(self) -> None:
payload = build_policy_payload(
name="Test Rule",
source_ip="1.2.3.4",
ip_version="IPV4",
src_zone_id="wan-zone",
dst_zone_id="lan-zone",
action_type="ALLOW",
)
assert payload["name"] == "Test Rule"
assert payload["source"]["zoneId"] == "wan-zone"
assert payload["destination"]["zoneId"] == "lan-zone"
assert payload["action"]["type"] == "ALLOW"
assert payload["action"]["allowReturnTraffic"] is True
assert payload["ipProtocolScope"]["ipVersion"] == "IPV4"
assert payload["enabled"] is True
assert payload["loggingEnabled"] is False
# Check source IP filter
tf = payload["source"]["trafficFilter"]
assert tf is not None
assert tf["type"] == "IP_ADDRESS"
assert tf["ipAddressFilter"]["type"] == "IP_ADDRESSES"
assert tf["ipAddressFilter"]["items"][0]["value"] == "1.2.3.4"
def test_with_dest_ports(self) -> None:
payload = build_policy_payload(
name="Test Rule",
source_ip="1.2.3.4",
ip_version="IPV4",
src_zone_id="wan-zone",
dst_zone_id="lan-zone",
action_type="ALLOW",
dest_ports=[22, 80, 443],
)
dest_tf = payload["destination"]["trafficFilter"]
assert dest_tf is not None
assert dest_tf["type"] == "PORT"
port_filter = dest_tf["portFilter"]
assert port_filter is not None
assert port_filter["type"] == "PORTS"
assert len(port_filter["items"]) == 3
assert port_filter["items"][0] == {"type": "PORT_NUMBER", "value": 22}
def test_with_port_ranges(self) -> None:
payload = build_policy_payload(
name="Test Rule",
source_ip="1.2.3.4",
ip_version="IPV4",
src_zone_id="wan-zone",
dst_zone_id="lan-zone",
action_type="ALLOW",
dest_port_ranges=[{"start": 8000, "stop": 9000}],
)
dest_tf = payload["destination"]["trafficFilter"]
assert dest_tf is not None
assert dest_tf["type"] == "PORT"
port_filter = dest_tf["portFilter"]
assert port_filter is not None
assert port_filter["items"][0] == {"type": "PORT_NUMBER_RANGE", "start": 8000, "stop": 9000}
def test_with_protocol(self) -> None:
payload = build_policy_payload(
name="Test Rule",
source_ip="1.2.3.4",
ip_version="IPV6",
src_zone_id="wan-zone",
dst_zone_id="lan-zone",
action_type="ALLOW",
protocol="tcp",
)
assert payload["ipProtocolScope"]["ipVersion"] == "IPV6"
pf = payload["ipProtocolScope"]["protocolFilter"]
assert pf is not None
assert pf["protocol"] == {"name": "TCP"}
def test_block_action(self) -> None:
payload = build_policy_payload(
name="Block Rule",
source_ip="5.6.7.8",
ip_version="IPV4",
src_zone_id="wan-zone",
dst_zone_id="lan-zone",
action_type="BLOCK",
)
assert payload["action"]["type"] == "BLOCK"
assert "allowReturnTraffic" not in payload["action"]
class TestGetSourceIpFromPolicy:
def test_extracts_ip_from_policy(self) -> None:
policy: dict[str, Any] = {
"id": "pol-1",
"name": "Test",
"source": {
"zoneId": "wan-zone",
"trafficFilter": {
"type": "IP_ADDRESS",
"ipAddressFilter": {
"type": "IP_ADDRESSES",
"items": [{"type": "IP_ADDRESS", "value": "1.2.3.4"}],
},
},
},
"destination": {"zoneId": "lan-zone", "trafficFilter": None},
"action": {"type": "ALLOW", "allowReturnTraffic": True},
"ipProtocolScope": {"ipVersion": "IPV4"},
"enabled": True,
"loggingEnabled": False,
"index": 0,
"metadata": {},
}
ip = get_source_ip_from_policy(policy) # type: ignore[arg-type]
assert ip == "1.2.3.4"
def test_returns_none_no_traffic_filter(self) -> None:
policy: dict[str, Any] = {
"source": {"zoneId": "wan-zone", "trafficFilter": None},
}
ip = get_source_ip_from_policy(policy) # type: ignore[arg-type]
assert ip is None
def test_returns_none_wrong_filter_type(self) -> None:
policy: dict[str, Any] = {
"source": {
"zoneId": "wan-zone",
"trafficFilter": {"type": "NETWORK"},
},
}
ip = get_source_ip_from_policy(policy) # type: ignore[arg-type]
assert ip is None
def test_returns_none_empty_items(self) -> None:
policy: dict[str, Any] = {
"source": {
"zoneId": "wan-zone",
"trafficFilter": {
"type": "IP_ADDRESS",
"ipAddressFilter": {"type": "IP_ADDRESSES", "items": []},
},
},
}
ip = get_source_ip_from_policy(policy) # type: ignore[arg-type]
assert ip is None