| """Unit tests for GraphSAGE model and graph embedding service. |
| |
| Tests the core FedGNN components: |
| 1. GraphSAGE model forward pass, embedding dimensions, and gradient flow |
| 2. Node feature extraction from entity dictionaries |
| 3. Model weight serialization/deserialization for federation |
| 4. Graph embedding service: local graph construction, training, similarity search |
| 5. Integration with FL engine: GNN weight aggregation validation |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import pytest |
| import torch |
|
|
| from app.application.services.graph_embedding_model import ( |
| NODE_FEATURE_DIM, |
| GraphSAGELayer, |
| GraphSAGEModel, |
| extract_node_features, |
| ) |
| from app.application.services.graph_embedding_service import GraphEmbeddingService |
| from app.application.services.graph_engine import GraphEngine |
| from app.domain.enums import AggregationMethod |
|
|
| |
|
|
|
|
| class TestNodeFeatureExtraction: |
| """Tests for extract_node_features().""" |
|
|
| def test_feature_dimension(self) -> None: |
| """Feature vector must be exactly NODE_FEATURE_DIM (12) dimensions.""" |
| entity = { |
| "entity_type": "customer", |
| "risk_level": "high", |
| "alert_count": 5, |
| } |
| features = extract_node_features(entity, degree=3) |
| assert features.shape == (NODE_FEATURE_DIM,) |
| assert features.dtype == np.float32 |
|
|
| def test_entity_type_one_hot(self) -> None: |
| """Entity type should be one-hot encoded in positions 0-6.""" |
| for etype, expected_idx in [ |
| ("customer", 0), |
| ("merchant", 1), |
| ("device", 2), |
| ("card", 3), |
| ]: |
| features = extract_node_features({"entity_type": etype}) |
| |
| assert features[expected_idx] == 1.0 |
| for i in range(7): |
| if i != expected_idx: |
| assert features[i] == 0.0 |
|
|
| def test_risk_level_ordinal(self) -> None: |
| """Risk level must be permanently blinded (0.0) and unmasked leakage must fail closed.""" |
| features_masked = extract_node_features({"risk_level": "critical"}) |
| assert features_masked[7] == 0.0, "Risk level ordinal must be masked to 0.0 by default" |
|
|
| with pytest.raises(ValueError, match="Unmasked label leakage in graph features is strictly prohibited"): |
| extract_node_features({"risk_level": "critical"}, mask_label_leakage=False) |
|
|
| def test_alert_count_normalized(self) -> None: |
| """Alert count should be log-normalized and bounded.""" |
| features_zero = extract_node_features({"alert_count": 0}) |
| features_high = extract_node_features({"alert_count": 100}) |
| assert features_zero[8] == 0.0 |
| assert 0.0 < features_high[8] <= 1.5 |
|
|
| def test_degree_centrality(self) -> None: |
| """Degree should be log-normalized in position 9.""" |
| features_iso = extract_node_features({}, degree=0) |
| features_hub = extract_node_features({}, degree=50) |
| assert features_iso[9] == 0.0 |
| assert features_hub[9] > features_iso[9] |
|
|
| def test_empty_entity_dict(self) -> None: |
| """Should handle empty dict without crashing.""" |
| features = extract_node_features({}) |
| assert features.shape == (NODE_FEATURE_DIM,) |
| |
| assert features[0] == 1.0 |
|
|
|
|
| |
|
|
|
|
| class TestGraphSAGELayer: |
| """Tests for individual GraphSAGE message-passing layers.""" |
|
|
| def test_output_shape(self) -> None: |
| """Layer output shape should be (N, out_dim).""" |
| layer = GraphSAGELayer(in_dim=12, out_dim=64) |
| features = torch.randn(5, 12) |
| adj = [[1, 2], [0, 2], [0, 1, 3], [2, 4], [3]] |
|
|
| output = layer(features, adj) |
| assert output.shape == (5, 64) |
|
|
| def test_isolated_node(self) -> None: |
| """Isolated nodes (no neighbors) should still produce embeddings.""" |
| layer = GraphSAGELayer(in_dim=12, out_dim=32) |
| features = torch.randn(3, 12) |
| adj = [[], [], []] |
|
|
| output = layer(features, adj) |
| assert output.shape == (3, 32) |
| |
| assert not torch.all(output == 0) |
|
|
| def test_l2_normalization(self) -> None: |
| """Output embeddings should be L2-normalized to unit sphere.""" |
| layer = GraphSAGELayer(in_dim=12, out_dim=32) |
| features = torch.randn(4, 12) |
| adj = [[1], [0, 2], [1, 3], [2]] |
|
|
| output = layer(features, adj) |
| norms = torch.norm(output, dim=1) |
| |
| for i in range(4): |
| if torch.any(output[i] != 0): |
| assert abs(norms[i].item() - 1.0) < 1e-5 |
|
|
| def test_gradient_flow(self) -> None: |
| """Gradients should flow through the layer.""" |
| layer = GraphSAGELayer(in_dim=12, out_dim=32) |
| features = torch.randn(3, 12, requires_grad=True) |
| adj = [[1, 2], [0], [0]] |
|
|
| output = layer(features, adj) |
| loss = output.sum() |
| loss.backward() |
|
|
| assert features.grad is not None |
| assert not torch.all(features.grad == 0) |
|
|
|
|
| |
|
|
|
|
| class TestGraphSAGEModel: |
| """Tests for the full GraphSAGE model.""" |
|
|
| def test_embedding_dimension(self) -> None: |
| """Model should produce embeddings of the configured dimension.""" |
| model = GraphSAGEModel(input_dim=12, hidden_dim=64, embedding_dim=32, num_layers=2) |
| features = torch.randn(5, 12) |
| adj = [[1, 2], [0, 3], [0, 4], [1], [2]] |
|
|
| embeddings = model.get_embeddings(features, adj) |
| assert embeddings.shape == (5, 32) |
|
|
| def test_forward_returns_embeddings_and_predictions(self) -> None: |
| """Forward pass should return both embeddings and fraud predictions.""" |
| model = GraphSAGEModel(input_dim=12, hidden_dim=64, embedding_dim=32, num_layers=2) |
| features = torch.randn(4, 12) |
| adj = [[1], [0, 2], [1, 3], [2]] |
|
|
| embeddings, predictions = model(features, adj) |
| assert embeddings.shape == (4, 32) |
| assert predictions.shape == (4,) |
| |
| assert torch.all(predictions >= 0) |
| assert torch.all(predictions <= 1) |
|
|
| def test_model_weights_serialization_roundtrip(self) -> None: |
| """Weights should survive serialize → deserialize without loss.""" |
| model = GraphSAGEModel(input_dim=12, hidden_dim=64, embedding_dim=32, num_layers=2) |
|
|
| |
| for param in model.parameters(): |
| param.data = torch.randn_like(param) |
|
|
| |
| weights = model.to_model_weights(include_classifier=True) |
| assert len(weights.flat_weights) > 0 |
| assert len(weights.layer_shapes) > 0 |
|
|
| |
| model2 = GraphSAGEModel(input_dim=12, hidden_dim=64, embedding_dim=32, num_layers=2) |
| model2.load_model_weights(weights, include_classifier=True) |
|
|
| |
| for p1, p2 in zip(model.parameters(), model2.parameters()): |
| assert torch.allclose(p1.data, p2.data, atol=1e-6) |
|
|
| def test_model_weights_parameter_count(self) -> None: |
| """Serialized weight count should match model parameter count.""" |
| model = GraphSAGEModel(input_dim=12, hidden_dim=128, embedding_dim=64, num_layers=2) |
| weights = model.to_model_weights(include_classifier=True) |
|
|
| expected_params = sum(p.numel() for p in model.parameters()) |
| assert weights.num_parameters == expected_params |
|
|
|
|
| |
|
|
|
|
| class TestGNNAggregation: |
| """Tests for GNN weight aggregation via fl_engine.""" |
|
|
| def test_fedavg_gnn_weights(self) -> None: |
| """FedAvg should correctly average GNN model weights.""" |
| from app.application.services.fl_engine import FederatedLearningEngine |
|
|
| |
| model_a = GraphSAGEModel(input_dim=12, hidden_dim=32, embedding_dim=16, num_layers=1) |
| model_b = GraphSAGEModel(input_dim=12, hidden_dim=32, embedding_dim=16, num_layers=1) |
|
|
| for p in model_a.parameters(): |
| p.data = torch.ones_like(p) * 2.0 |
| for p in model_b.parameters(): |
| p.data = torch.ones_like(p) * 4.0 |
|
|
| weights_a = model_a.to_model_weights() |
| weights_b = model_b.to_model_weights() |
|
|
| |
| engine = FederatedLearningEngine.__new__(FederatedLearningEngine) |
|
|
| aggregated = engine.aggregate_parameters( |
| [weights_a, weights_b], |
| [100, 100], |
| method=AggregationMethod.FED_AVG, |
| ) |
|
|
| |
| for w in aggregated.flat_weights: |
| assert abs(w - 3.0) < 1e-6 |
|
|
| def test_gnn_layer_shape_validation(self) -> None: |
| """aggregate_graph_parameters should reject mismatched layer shapes.""" |
| from app.application.services.fl_engine import FederatedLearningEngine |
|
|
| |
| model_a = GraphSAGEModel(input_dim=12, hidden_dim=32, embedding_dim=16, num_layers=2) |
| model_b = GraphSAGEModel(input_dim=12, hidden_dim=64, embedding_dim=16, num_layers=2) |
|
|
| weights_a = model_a.to_model_weights() |
| weights_b = model_b.to_model_weights() |
|
|
| engine = FederatedLearningEngine.__new__(FederatedLearningEngine) |
|
|
| with pytest.raises(ValueError, match="GNN layer shape mismatch"): |
| engine.aggregate_graph_parameters( |
| [weights_a, weights_b], |
| [100, 100], |
| ) |
|
|
|
|
| |
|
|
|
|
| class TestGraphEmbeddingService: |
| """Tests for the graph embedding orchestration service.""" |
|
|
| def test_empty_graph_training(self) -> None: |
| """Training on an empty graph should return zero loss without crashing.""" |
| service = GraphEmbeddingService(graph_engine=GraphEngine()) |
| weights, metrics = service.train_local_gnn(bank_id="bank_nonexistent") |
| assert metrics["loss"] == 0.0 |
| assert metrics["num_nodes"] == 0 |
|
|
| def test_embedding_stats_empty(self) -> None: |
| """Stats should return zeros when no embeddings exist.""" |
| service = GraphEmbeddingService() |
| stats = service.get_embedding_stats() |
| assert stats["num_embedded_nodes"] == 0 |
|
|
| def test_find_similar_no_embeddings(self) -> None: |
| """Similarity search should return empty when no embeddings cached.""" |
| service = GraphEmbeddingService() |
| results = service.find_similar_entities("nonexistent") |
| assert results == [] |
|
|
| def test_get_all_embeddings_serializable(self) -> None: |
| """get_all_embeddings should return JSON-serializable dict.""" |
| service = GraphEmbeddingService() |
| |
| service._embeddings["test_node"] = np.array([0.1, 0.2, 0.3], dtype=np.float32) |
|
|
| result = service.get_all_embeddings() |
| assert "test_node" in result |
| assert isinstance(result["test_node"], list) |
| assert len(result["test_node"]) == 3 |
|
|
| def test_cosine_similarity_search(self) -> None: |
| """Cosine similarity search should find similar embeddings.""" |
| service = GraphEmbeddingService() |
|
|
| |
| service._embeddings["node_a"] = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) |
| service._embeddings["node_b"] = np.array([0.9, 0.1, 0.0, 0.0], dtype=np.float32) |
| service._embeddings["node_c"] = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32) |
|
|
| results = service.find_similar_entities("node_a", top_k=5, threshold=0.5) |
|
|
| |
| entity_ids = [r["entity_id"] for r in results] |
| assert "node_b" in entity_ids |
|
|
| |
| assert "node_c" not in entity_ids |
|
|
|
|
| class TestGraphSAGETensorEdgeIndex: |
| """Tests evaluating PyG-style edge_index tensors and masked semi-supervised loss.""" |
|
|
| def test_graphsage_layer_forward_with_edge_index(self) -> None: |
| """Verify GraphSAGELayer forward pass with tensor edge_index produces valid embeddings.""" |
| torch.manual_seed(42) |
| layer = GraphSAGELayer(in_dim=16, out_dim=8) |
| node_features = torch.randn(10, 16) |
| edge_index = torch.tensor([[0, 1, 2, 3, 4], [1, 2, 3, 4, 0]], dtype=torch.long) |
|
|
| out = layer(node_features, edge_index=edge_index) |
| assert out.shape == (10, 8) |
| |
| for i in range(10): |
| if torch.any(out[i] != 0): |
| norm = torch.norm(out[i], p=2) |
| assert abs(norm.item() - 1.0) < 1e-5 |
|
|
| def test_graphsage_model_edge_index_and_166_features(self) -> None: |
| """Verify GraphSAGEModel handles 166-dim Elliptic features with edge_index.""" |
| model = GraphSAGEModel(input_dim=166, hidden_dim=32, embedding_dim=16, num_layers=2) |
| node_features = torch.randn(20, 166) |
| edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype=torch.long) |
|
|
| embeddings, predictions = model(node_features, edge_index=edge_index) |
| assert embeddings.shape == (20, 16) |
| assert predictions.shape == (20,) |
| assert (predictions >= 0.0).all() and (predictions <= 1.0).all() |
|
|
| def test_graphsage_model_compute_loss_with_mask(self) -> None: |
| """Verify masked loss computation strictly isolates labeled nodes.""" |
| model = GraphSAGEModel(input_dim=16, hidden_dim=16, embedding_dim=8) |
| predictions = torch.tensor([0.9, 0.1, 0.8, 0.2], requires_grad=True) |
| targets = torch.tensor([1, 0, 0, 1]) |
| mask = torch.tensor([True, True, False, False]) |
|
|
| loss = model.compute_loss(predictions, targets, mask=mask) |
| assert loss.item() > 0.0 |
|
|
| |
| loss.backward() |
| assert predictions.grad is not None |
| assert predictions.grad[0] != 0.0 |
| assert predictions.grad[1] != 0.0 |
| assert predictions.grad[2] == 0.0 |
| assert predictions.grad[3] == 0.0 |
|
|
|
|
