206 lines
7.5 KiB
Python
206 lines
7.5 KiB
Python
"""Tests for heatmap_service — pure helpers, error paths, and dispatch."""
|
|
|
|
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from app.services.heatmap_service import (
|
|
_score_to_color,
|
|
_build_layer_skeleton,
|
|
_parse_csv,
|
|
_format_tactic,
|
|
build_navigator_export,
|
|
build_threat_actor_layer,
|
|
build_campaign_layer,
|
|
SUPPORTED_LAYER_TYPES,
|
|
ATTACK_VERSION,
|
|
NAVIGATOR_VERSION,
|
|
LAYER_VERSION,
|
|
DOMAIN,
|
|
)
|
|
from app.domain.errors import BusinessRuleViolation, EntityNotFoundError
|
|
|
|
|
|
# ── Pure helpers ──────────────────────────────────────────────────────
|
|
|
|
|
|
class TestScoreToColor:
|
|
def test_zero_returns_grey(self):
|
|
assert _score_to_color(0) == "#d3d3d3"
|
|
|
|
def test_negative_returns_grey(self):
|
|
assert _score_to_color(-10) == "#d3d3d3"
|
|
|
|
def test_low_returns_red(self):
|
|
assert _score_to_color(25) == "#ff6666"
|
|
|
|
def test_medium_returns_orange(self):
|
|
assert _score_to_color(50) == "#ff9933"
|
|
|
|
def test_high_returns_yellow(self):
|
|
assert _score_to_color(75) == "#ffff66"
|
|
|
|
def test_max_returns_green(self):
|
|
assert _score_to_color(100) == "#66ff66"
|
|
|
|
|
|
class TestBuildLayerSkeleton:
|
|
def test_has_required_keys(self):
|
|
layer = _build_layer_skeleton("Test Layer", "A description")
|
|
assert layer["name"] == "Test Layer"
|
|
assert layer["description"] == "A description"
|
|
assert layer["domain"] == DOMAIN
|
|
assert layer["techniques"] == []
|
|
assert layer["versions"]["attack"] == ATTACK_VERSION
|
|
assert layer["versions"]["navigator"] == NAVIGATOR_VERSION
|
|
assert layer["versions"]["layer"] == LAYER_VERSION
|
|
|
|
def test_default_gradient(self):
|
|
layer = _build_layer_skeleton("X", "Y")
|
|
assert layer["gradient"]["minValue"] == 0
|
|
assert layer["gradient"]["maxValue"] == 100
|
|
assert len(layer["gradient"]["colors"]) == 3
|
|
|
|
def test_custom_gradient(self):
|
|
layer = _build_layer_skeleton("X", "Y", gradient_colors=["#000", "#fff"])
|
|
assert layer["gradient"]["colors"] == ["#000", "#fff"]
|
|
|
|
|
|
class TestParseCsv:
|
|
def test_none_returns_none(self):
|
|
assert _parse_csv(None) is None
|
|
|
|
def test_empty_string_returns_none(self):
|
|
assert _parse_csv("") is None
|
|
|
|
def test_single_value(self):
|
|
assert _parse_csv("windows") == ["windows"]
|
|
|
|
def test_multiple_values_with_spaces(self):
|
|
assert _parse_csv("windows, linux, macos") == ["windows", "linux", "macos"]
|
|
|
|
def test_empty_elements_filtered(self):
|
|
assert _parse_csv("a,,b") == ["a", "b"]
|
|
|
|
|
|
class TestFormatTactic:
|
|
def test_none_returns_empty(self):
|
|
assert _format_tactic(None) == ""
|
|
|
|
def test_empty_returns_empty(self):
|
|
assert _format_tactic("") == ""
|
|
|
|
def test_lowercases(self):
|
|
assert _format_tactic("Initial Access") == "initial access"
|
|
|
|
def test_comma_separated_takes_first(self):
|
|
assert _format_tactic("Execution, Persistence") == "execution"
|
|
|
|
|
|
# ── build_navigator_export dispatch ───────────────────────────────────
|
|
|
|
|
|
def _mock_db():
|
|
return MagicMock()
|
|
|
|
|
|
class TestBuildNavigatorExport:
|
|
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
|
|
|
|
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
|
|
|
|
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
|
|
|
|
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"):
|
|
build_navigator_export(_mock_db(), "nonexistent")
|
|
|
|
def test_missing_layer_id_for_threat_actor(self):
|
|
with pytest.raises(BusinessRuleViolation, match="layer_id is required"):
|
|
build_navigator_export(_mock_db(), "threat-actor")
|
|
|
|
def test_missing_layer_id_for_campaign(self):
|
|
with pytest.raises(BusinessRuleViolation, match="layer_id is required"):
|
|
build_navigator_export(_mock_db(), "campaign")
|
|
|
|
def test_supported_layer_types_complete(self):
|
|
assert SUPPORTED_LAYER_TYPES == {
|
|
"coverage", "detection-rules", "threat-actor", "campaign",
|
|
}
|
|
|
|
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 ───────────────────────────────────────────
|
|
|
|
|
|
class TestEntityNotFound:
|
|
def _db_returning_none(self):
|
|
db = MagicMock()
|
|
db.query.return_value.filter.return_value.first.return_value = None
|
|
return db
|
|
|
|
def test_threat_actor_not_found(self):
|
|
with pytest.raises(EntityNotFoundError, match="ThreatActor"):
|
|
build_threat_actor_layer(self._db_returning_none(), "bad-id")
|
|
|
|
def test_campaign_not_found(self):
|
|
with pytest.raises(EntityNotFoundError, match="Campaign"):
|
|
build_campaign_layer(self._db_returning_none(), "bad-id")
|