291 lines
9.5 KiB
Python
291 lines
9.5 KiB
Python
"""Unit tests for v3 correlation-aware clustering.
|
||
|
||
Tests for compute_n_eff, compute_cluster_llr, and cluster_evidence functions.
|
||
|
||
Requirements validated: 4.1–4.5
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from datetime import datetime, timezone
|
||
|
||
import pytest
|
||
|
||
from services.aggregation.scoring import EvidenceUnit
|
||
from services.aggregation.worker import (
|
||
cluster_evidence,
|
||
compute_cluster_llr,
|
||
compute_n_eff,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_unit(cluster_id: str = "cluster_a", symbol: str = "AAPL") -> EvidenceUnit:
|
||
"""Create a minimal EvidenceUnit for testing."""
|
||
return EvidenceUnit(
|
||
symbol=symbol,
|
||
layer="company",
|
||
event_type="earnings",
|
||
source_id="doc_1",
|
||
source_group="reuters",
|
||
timestamp=datetime(2024, 1, 15, 12, 0, tzinfo=timezone.utc),
|
||
horizon="7d",
|
||
direction=1,
|
||
sentiment_strength=0.8,
|
||
impact=0.7,
|
||
extraction_conf=0.9,
|
||
source_cred=0.85,
|
||
novelty=0.6,
|
||
event_base_rate=0.25,
|
||
cluster_id=cluster_id,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test: 3 identical articles from same source → n_eff < 3
|
||
# Requirement: 4.2, 4.3
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestNEffIdenticalArticles:
|
||
"""3 identical articles from same wire (default rho=0.80) → n_eff < 3."""
|
||
|
||
def test_n_eff_less_than_cluster_size(self):
|
||
llrs = [1.0, 1.0, 1.0]
|
||
# Default correlations: rho=0.80 for all pairs (same wire/source)
|
||
n_eff = compute_n_eff(llrs)
|
||
|
||
# Formula: (3)² / (3 + 2×3×0.80×1×1) = 9 / (3 + 4.8) = 9/7.8 ≈ 1.154
|
||
expected = 9.0 / 7.8
|
||
assert n_eff < 3.0
|
||
assert n_eff == pytest.approx(expected, rel=1e-6)
|
||
|
||
def test_n_eff_discounts_correlated_signals(self):
|
||
"""Higher correlation → lower n_eff."""
|
||
llrs = [1.0, 1.0, 1.0]
|
||
n_eff_correlated = compute_n_eff(llrs) # default rho=0.80
|
||
|
||
# Independent: rho=0.0
|
||
identity = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
|
||
n_eff_independent = compute_n_eff(llrs, correlations=identity)
|
||
|
||
assert n_eff_correlated < n_eff_independent
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test: 3 independent articles → n_eff ≈ 3
|
||
# Requirement: 4.2, 4.3
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestNEffIndependentArticles:
|
||
"""3 independent articles (rho=0.0) → n_eff = 3.0."""
|
||
|
||
def test_n_eff_equals_cluster_size(self):
|
||
llrs = [1.0, 1.0, 1.0]
|
||
# Zero off-diagonal correlations
|
||
correlations = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
|
||
n_eff = compute_n_eff(llrs, correlations=correlations)
|
||
|
||
# n_eff = (3)² / (3 + 0) = 3.0
|
||
assert n_eff == pytest.approx(3.0, rel=1e-6)
|
||
|
||
def test_n_eff_with_varying_magnitudes(self):
|
||
"""Independent signals with different magnitudes still give n <= cluster size."""
|
||
llrs = [0.5, 1.0, 2.0]
|
||
correlations = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
|
||
n_eff = compute_n_eff(llrs, correlations=correlations)
|
||
|
||
# With zero correlations, n_eff = (sum |w|)^2 / sum(w^2)
|
||
# = (0.5+1.0+2.0)^2 / (0.25+1.0+4.0) = 12.25 / 5.25 ≈ 2.333
|
||
expected = (3.5**2) / (0.25 + 1.0 + 4.0)
|
||
assert n_eff == pytest.approx(expected, rel=1e-6)
|
||
assert n_eff <= 3.0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test: Single signal cluster → n_eff = 1.0
|
||
# Requirement: 4.2
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestNEffSingleSignal:
|
||
"""Single signal in a cluster → n_eff = 1.0."""
|
||
|
||
def test_single_signal(self):
|
||
llrs = [0.5]
|
||
n_eff = compute_n_eff(llrs)
|
||
assert n_eff == 1.0
|
||
|
||
def test_empty_cluster(self):
|
||
llrs: list[float] = []
|
||
n_eff = compute_n_eff(llrs)
|
||
assert n_eff == 1.0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test: Cluster LLR clamp at ±2.5
|
||
# Requirement: 4.4, 4.5
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestClusterLLRClamp:
|
||
"""Cluster LLR is clamped to [-2.5, 2.5]."""
|
||
|
||
def test_positive_clamp(self):
|
||
"""Very large positive LLRs with high n_eff → clamped to 2.5."""
|
||
llrs = [2.0, 2.0, 2.0, 2.0, 2.0]
|
||
# Use independent correlations for max n_eff
|
||
correlations = [
|
||
[1, 0, 0, 0, 0],
|
||
[0, 1, 0, 0, 0],
|
||
[0, 0, 1, 0, 0],
|
||
[0, 0, 0, 1, 0],
|
||
[0, 0, 0, 0, 1],
|
||
]
|
||
n_eff = compute_n_eff(llrs, correlations=correlations)
|
||
cluster_llr = compute_cluster_llr(llrs, n_eff)
|
||
|
||
# weighted_mean = 2.0, sqrt(5) ≈ 2.236, raw = 4.47 → clamp to 2.5
|
||
assert cluster_llr == pytest.approx(2.5, rel=1e-6)
|
||
|
||
def test_negative_clamp(self):
|
||
"""Very negative LLRs with high n_eff → clamped to -2.5."""
|
||
llrs = [-2.0, -2.0, -2.0, -2.0, -2.0]
|
||
correlations = [
|
||
[1, 0, 0, 0, 0],
|
||
[0, 1, 0, 0, 0],
|
||
[0, 0, 1, 0, 0],
|
||
[0, 0, 0, 1, 0],
|
||
[0, 0, 0, 0, 1],
|
||
]
|
||
n_eff = compute_n_eff(llrs, correlations=correlations)
|
||
cluster_llr = compute_cluster_llr(llrs, n_eff)
|
||
|
||
assert cluster_llr == pytest.approx(-2.5, rel=1e-6)
|
||
|
||
def test_within_bounds_no_clamp(self):
|
||
"""Small LLRs with low n_eff → no clamping needed."""
|
||
llrs = [0.3, 0.4]
|
||
n_eff = compute_n_eff(llrs)
|
||
cluster_llr = compute_cluster_llr(llrs, n_eff)
|
||
|
||
assert -2.5 <= cluster_llr <= 2.5
|
||
# Should NOT be at the clamp boundary
|
||
assert abs(cluster_llr) < 2.5
|
||
|
||
def test_single_signal_clamp(self):
|
||
"""Single signal beyond clamp → clamped."""
|
||
llrs = [3.0]
|
||
cluster_llr = compute_cluster_llr(llrs, n_eff=1.0)
|
||
assert cluster_llr == pytest.approx(2.5, rel=1e-6)
|
||
|
||
def test_single_signal_negative_clamp(self):
|
||
"""Single negative signal beyond clamp → clamped to -2.5."""
|
||
llrs = [-3.0]
|
||
cluster_llr = compute_cluster_llr(llrs, n_eff=1.0)
|
||
assert cluster_llr == pytest.approx(-2.5, rel=1e-6)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test: All-zero LLRs → cluster_llr = 0.0
|
||
# Requirement: 4.4
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestClusterLLRZero:
|
||
"""All-zero LLRs produce zero cluster LLR."""
|
||
|
||
def test_all_zeros(self):
|
||
llrs = [0.0, 0.0, 0.0]
|
||
n_eff = compute_n_eff(llrs)
|
||
cluster_llr = compute_cluster_llr(llrs, n_eff)
|
||
assert cluster_llr == 0.0
|
||
|
||
def test_empty_llrs(self):
|
||
"""Empty LLR list → 0.0."""
|
||
cluster_llr = compute_cluster_llr([], n_eff=1.0)
|
||
assert cluster_llr == 0.0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test: Grouping by correct key dimensions (cluster_id)
|
||
# Requirement: 4.1
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestClusterEvidence:
|
||
"""cluster_evidence groups EvidenceUnits by cluster_id."""
|
||
|
||
def test_grouping_by_cluster_id(self):
|
||
"""Units with same cluster_id are grouped together."""
|
||
unit_a1 = _make_unit(cluster_id="cluster_a")
|
||
unit_a2 = _make_unit(cluster_id="cluster_a")
|
||
unit_b1 = _make_unit(cluster_id="cluster_b")
|
||
|
||
units = [unit_a1, unit_a2, unit_b1]
|
||
llrs = [1.0, 0.5, -0.3]
|
||
|
||
clusters = cluster_evidence(units, llrs)
|
||
|
||
assert len(clusters) == 2
|
||
|
||
# Find clusters by id
|
||
cluster_map = {c.cluster_id: c for c in clusters}
|
||
|
||
assert "cluster_a" in cluster_map
|
||
assert "cluster_b" in cluster_map
|
||
|
||
# Cluster A has 2 units
|
||
assert len(cluster_map["cluster_a"].units) == 2
|
||
assert cluster_map["cluster_a"].llrs == [1.0, 0.5]
|
||
|
||
# Cluster B has 1 unit
|
||
assert len(cluster_map["cluster_b"].units) == 1
|
||
assert cluster_map["cluster_b"].llrs == [-0.3]
|
||
|
||
def test_single_cluster(self):
|
||
"""All units with same cluster_id → one cluster."""
|
||
units = [_make_unit(cluster_id="only") for _ in range(4)]
|
||
llrs = [0.1, 0.2, 0.3, 0.4]
|
||
|
||
clusters = cluster_evidence(units, llrs)
|
||
|
||
assert len(clusters) == 1
|
||
assert clusters[0].cluster_id == "only"
|
||
assert len(clusters[0].units) == 4
|
||
assert clusters[0].llrs == [0.1, 0.2, 0.3, 0.4]
|
||
|
||
def test_each_unit_different_cluster(self):
|
||
"""Each unit in its own cluster → N clusters."""
|
||
units = [_make_unit(cluster_id=f"c_{i}") for i in range(5)]
|
||
llrs = [0.1 * i for i in range(5)]
|
||
|
||
clusters = cluster_evidence(units, llrs)
|
||
|
||
assert len(clusters) == 5
|
||
for c in clusters:
|
||
assert len(c.units) == 1
|
||
|
||
def test_empty_input(self):
|
||
"""No units → no clusters."""
|
||
clusters = cluster_evidence([], [])
|
||
assert clusters == []
|
||
|
||
def test_llrs_parallel_to_units(self):
|
||
"""LLRs are correctly associated with their units."""
|
||
unit_x = _make_unit(cluster_id="x")
|
||
unit_y = _make_unit(cluster_id="y")
|
||
unit_x2 = _make_unit(cluster_id="x")
|
||
|
||
units = [unit_x, unit_y, unit_x2]
|
||
llrs = [1.5, -0.7, 2.3]
|
||
|
||
clusters = cluster_evidence(units, llrs)
|
||
cluster_map = {c.cluster_id: c for c in clusters}
|
||
|
||
assert cluster_map["x"].llrs == [1.5, 2.3]
|
||
assert cluster_map["y"].llrs == [-0.7]
|