feat: make heatmap layers extensible via LayerRegistry (OCP)
This commit is contained in:
58
backend/tests/test_heatmap_layer_registry.py
Normal file
58
backend/tests/test_heatmap_layer_registry.py
Normal file
@@ -0,0 +1,58 @@
|
||||
"""Tests for the heatmap layer registry (OCP extensibility)."""
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services.heatmap_service import (
|
||||
LAYER_REGISTRY,
|
||||
SUPPORTED_LAYER_TYPES,
|
||||
register_layer,
|
||||
)
|
||||
|
||||
|
||||
def test_builtin_layer_types_registered():
|
||||
assert "coverage" in SUPPORTED_LAYER_TYPES
|
||||
assert "detection-rules" in SUPPORTED_LAYER_TYPES
|
||||
assert "threat-actor" in SUPPORTED_LAYER_TYPES
|
||||
assert "campaign" in SUPPORTED_LAYER_TYPES
|
||||
|
||||
|
||||
def test_register_custom_simple_layer():
|
||||
def my_layer(db, *, platforms=None, tactics=None, min_score=0):
|
||||
return {"name": "custom", "techniques": []}
|
||||
|
||||
register_layer("custom-test-layer", my_layer)
|
||||
assert "custom-test-layer" in LAYER_REGISTRY.supported_types
|
||||
|
||||
result = LAYER_REGISTRY.build(
|
||||
None, "custom-test-layer",
|
||||
platforms=None, tactics=None, min_score=0,
|
||||
)
|
||||
assert result["name"] == "custom"
|
||||
|
||||
|
||||
def test_register_custom_id_layer():
|
||||
def my_id_layer(db, layer_id, *, platforms=None, tactics=None, min_score=0):
|
||||
return {"name": f"entity-{layer_id}", "techniques": []}
|
||||
|
||||
register_layer("custom-id-layer", my_id_layer, requires_id=True)
|
||||
assert "custom-id-layer" in LAYER_REGISTRY.supported_types
|
||||
|
||||
result = LAYER_REGISTRY.build(
|
||||
None, "custom-id-layer",
|
||||
layer_id="abc-123", platforms=None, tactics=None, min_score=0,
|
||||
)
|
||||
assert result["name"] == "entity-abc-123"
|
||||
|
||||
|
||||
def test_unknown_layer_raises():
|
||||
from app.domain.errors import BusinessRuleViolation
|
||||
|
||||
with pytest.raises(BusinessRuleViolation, match="Unknown layer type"):
|
||||
LAYER_REGISTRY.build(None, "nonexistent-layer")
|
||||
|
||||
|
||||
def test_id_layer_without_id_raises():
|
||||
from app.domain.errors import BusinessRuleViolation
|
||||
|
||||
with pytest.raises(BusinessRuleViolation, match="layer_id is required"):
|
||||
LAYER_REGISTRY.build(None, "threat-actor", layer_id=None)
|
||||
@@ -104,33 +104,53 @@ def _mock_db():
|
||||
|
||||
|
||||
class TestBuildNavigatorExport:
|
||||
@patch("app.services.heatmap_service.build_coverage_layer")
|
||||
def test_dispatches_coverage(self, mock_build):
|
||||
mock_build.return_value = {"name": "coverage"}
|
||||
result = build_navigator_export(_mock_db(), "coverage")
|
||||
assert result["name"] == "coverage"
|
||||
mock_build.assert_called_once()
|
||||
def test_dispatches_coverage(self):
|
||||
from app.services.heatmap_service import LAYER_REGISTRY
|
||||
mock_build = MagicMock(return_value={"name": "coverage"})
|
||||
orig = LAYER_REGISTRY._simple["coverage"]
|
||||
LAYER_REGISTRY._simple["coverage"] = mock_build
|
||||
try:
|
||||
result = build_navigator_export(_mock_db(), "coverage")
|
||||
assert result["name"] == "coverage"
|
||||
mock_build.assert_called_once()
|
||||
finally:
|
||||
LAYER_REGISTRY._simple["coverage"] = orig
|
||||
|
||||
@patch("app.services.heatmap_service.build_detection_rules_layer")
|
||||
def test_dispatches_detection_rules(self, mock_build):
|
||||
mock_build.return_value = {"name": "rules"}
|
||||
result = build_navigator_export(_mock_db(), "detection-rules")
|
||||
assert result["name"] == "rules"
|
||||
mock_build.assert_called_once()
|
||||
def test_dispatches_detection_rules(self):
|
||||
from app.services.heatmap_service import LAYER_REGISTRY
|
||||
mock_build = MagicMock(return_value={"name": "rules"})
|
||||
orig = LAYER_REGISTRY._simple["detection-rules"]
|
||||
LAYER_REGISTRY._simple["detection-rules"] = mock_build
|
||||
try:
|
||||
result = build_navigator_export(_mock_db(), "detection-rules")
|
||||
assert result["name"] == "rules"
|
||||
mock_build.assert_called_once()
|
||||
finally:
|
||||
LAYER_REGISTRY._simple["detection-rules"] = orig
|
||||
|
||||
@patch("app.services.heatmap_service.build_threat_actor_layer")
|
||||
def test_dispatches_threat_actor_with_id(self, mock_build):
|
||||
mock_build.return_value = {"name": "actor"}
|
||||
result = build_navigator_export(_mock_db(), "threat-actor", layer_id="abc")
|
||||
assert result["name"] == "actor"
|
||||
mock_build.assert_called_once()
|
||||
def test_dispatches_threat_actor_with_id(self):
|
||||
from app.services.heatmap_service import LAYER_REGISTRY
|
||||
mock_build = MagicMock(return_value={"name": "actor"})
|
||||
orig = LAYER_REGISTRY._with_id["threat-actor"]
|
||||
LAYER_REGISTRY._with_id["threat-actor"] = mock_build
|
||||
try:
|
||||
result = build_navigator_export(_mock_db(), "threat-actor", layer_id="abc")
|
||||
assert result["name"] == "actor"
|
||||
mock_build.assert_called_once()
|
||||
finally:
|
||||
LAYER_REGISTRY._with_id["threat-actor"] = orig
|
||||
|
||||
@patch("app.services.heatmap_service.build_campaign_layer")
|
||||
def test_dispatches_campaign_with_id(self, mock_build):
|
||||
mock_build.return_value = {"name": "campaign"}
|
||||
result = build_navigator_export(_mock_db(), "campaign", layer_id="xyz")
|
||||
assert result["name"] == "campaign"
|
||||
mock_build.assert_called_once()
|
||||
def test_dispatches_campaign_with_id(self):
|
||||
from app.services.heatmap_service import LAYER_REGISTRY
|
||||
mock_build = MagicMock(return_value={"name": "campaign"})
|
||||
orig = LAYER_REGISTRY._with_id["campaign"]
|
||||
LAYER_REGISTRY._with_id["campaign"] = mock_build
|
||||
try:
|
||||
result = build_navigator_export(_mock_db(), "campaign", layer_id="xyz")
|
||||
assert result["name"] == "campaign"
|
||||
mock_build.assert_called_once()
|
||||
finally:
|
||||
LAYER_REGISTRY._with_id["campaign"] = orig
|
||||
|
||||
def test_unknown_layer_raises(self):
|
||||
with pytest.raises(BusinessRuleViolation, match="Unknown layer type"):
|
||||
@@ -149,17 +169,22 @@ class TestBuildNavigatorExport:
|
||||
"coverage", "detection-rules", "threat-actor", "campaign",
|
||||
}
|
||||
|
||||
@patch("app.services.heatmap_service.build_coverage_layer")
|
||||
def test_passes_filter_kwargs(self, mock_build):
|
||||
mock_build.return_value = {}
|
||||
build_navigator_export(
|
||||
_mock_db(), "coverage",
|
||||
platforms="windows", tactics="execution", min_score=50,
|
||||
)
|
||||
_, kwargs = mock_build.call_args
|
||||
assert kwargs["platforms"] == "windows"
|
||||
assert kwargs["tactics"] == "execution"
|
||||
assert kwargs["min_score"] == 50
|
||||
def test_passes_filter_kwargs(self):
|
||||
from app.services.heatmap_service import LAYER_REGISTRY
|
||||
mock_build = MagicMock(return_value={})
|
||||
orig = LAYER_REGISTRY._simple["coverage"]
|
||||
LAYER_REGISTRY._simple["coverage"] = mock_build
|
||||
try:
|
||||
build_navigator_export(
|
||||
_mock_db(), "coverage",
|
||||
platforms="windows", tactics="execution", min_score=50,
|
||||
)
|
||||
_, kwargs = mock_build.call_args
|
||||
assert kwargs["platforms"] == "windows"
|
||||
assert kwargs["tactics"] == "execution"
|
||||
assert kwargs["min_score"] == 50
|
||||
finally:
|
||||
LAYER_REGISTRY._simple["coverage"] = orig
|
||||
|
||||
|
||||
# ── Entity-not-found errors ───────────────────────────────────────────
|
||||
|
||||
Reference in New Issue
Block a user