59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
"""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)
|