"""Unit tests for GradientCompressionEngine and FederatedCommunicationProfiler. Tests cover: - Top-K gradient sparsification (accuracy, sparsity invariants, thresholding, edge cases) - Lossless compression and decompression roundtrips (zlib / zstd format, bit-exact recovery) - Binary quantization (FP32, FP16 half-precision, INT8 symmetric affine) - Coordinate sparse (COO) encoding and decoding - Wire payload transfer size hierarchy and compression profiling - Protocol communication profiling (FedAvg vs FedProx vs SCAFFOLD vs Curve25519 SecAgg vs TenSEAL CKKS) - Integration with FederatedLearningEngine communication telemetry """ from __future__ import annotations import math from unittest.mock import MagicMock import numpy as np import pytest from app.application.services.fl_engine import FederatedLearningEngine from app.domain.enums import AggregationMethod from app.infrastructure.security.compression_engine import ( BandwidthBenchmarkSuiteResult, FederatedCommunicationProfiler, FederatedProtocol, GradientCompressionEngine, ProtocolCommunicationProfile, SerializationFormat, ) class TestTopKSparsification: """Validates Top-K gradient sparsification algorithms and boundary behavior.""" @pytest.fixture def engine(self) -> GradientCompressionEngine: return GradientCompressionEngine(default_k_percent=0.20) def test_sparsify_top_k_retains_exact_percentage(self, engine: GradientCompressionEngine) -> None: """K=20% on 1,000 items retains exactly 200 non-zero elements.""" weights = [float(i) for i in range(1, 1001)] sparsified = engine.sparsify_top_k(weights, k_percent=0.20) assert len(sparsified) == 1000 non_zero = [w for w in sparsified if w != 0.0] assert len(non_zero) == 200 # The retained 200 should be 801 to 1000 assert min(non_zero) == 801.0 assert max(non_zero) == 1000.0 def test_sparsify_top_k_preserves_large_negative_magnitudes( self, engine: GradientCompressionEngine ) -> None: """Top-K operates on absolute magnitudes, preserving large negative gradients.""" weights = [-100.0, 0.5, -0.2, 50.0, 0.1, -80.0, 0.01, 20.0, -0.05, 0.02] # Keep top 30% (3 items): -100.0, -80.0, 50.0 sparsified = engine.sparsify_top_k(weights, k_percent=0.30) retained = {w for w in sparsified if w != 0.0} assert retained == {-100.0, -80.0, 50.0} def test_sparsify_top_k_clamping_and_defaults(self, engine: GradientCompressionEngine) -> None: """k_percent is clamped between 0.01 and 1.0.""" weights = [1.0, 2.0, 3.0, 4.0, 5.0] # None uses default_k_percent (0.20 -> 1 element) default_res = engine.sparsify_top_k(weights, k_percent=None) assert sum(1 for w in default_res if w != 0.0) == 1 # 0.0 clamped to 0.01 (at least 1 element) zero_res = engine.sparsify_top_k(weights, k_percent=0.0) assert sum(1 for w in zero_res if w != 0.0) == 1 # 1.5 clamped to 1.0 (all 5 elements retained) over_res = engine.sparsify_top_k(weights, k_percent=1.5) assert sum(1 for w in over_res if w != 0.0) == 5 def test_sparsify_top_k_empty_and_zero_inputs(self, engine: GradientCompressionEngine) -> None: """Handles empty and all-zero weight vectors safely.""" assert engine.sparsify_top_k([]) == [] assert engine.sparsify_top_k(None) == [] # type: ignore[arg-type] all_zeros = [0.0] * 50 res = engine.sparsify_top_k(all_zeros, k_percent=0.20) assert len(res) == 50 assert all(w == 0.0 for w in res) class TestLosslessCompression: """Validates lossless byte compression and decompression.""" @pytest.fixture def engine(self) -> GradientCompressionEngine: return GradientCompressionEngine(compression_level=6) def test_compress_decompress_roundtrip(self, engine: GradientCompressionEngine) -> None: """Compressed payload decompresses to identical bytes.""" original = b"FEDERATED_LEARNING_BANK_GRADIENT_UPDATE_PAYLOAD_TEST_VECTOR_12345" * 20 compressed = engine.compress_payload(original) assert len(compressed) < len(original) decompressed = engine.decompress_payload(compressed) assert decompressed == original def test_compress_empty_payload(self, engine: GradientCompressionEngine) -> None: """Empty input bytes returns empty bytes.""" assert engine.compress_payload(b"") == b"" assert engine.decompress_payload(b"") == b"" def test_compress_gradient_vector_backward_compatibility( self, engine: GradientCompressionEngine ) -> None: """Maintains backward compatibility with legacy compress_gradient_vector.""" weights = [0.1 * i for i in range(100)] meta = engine.compress_gradient_vector(weights, k_percent=0.10) assert meta["original_length"] == 100 assert meta["sparsified_length"] == 100 assert meta["compressed_size"] > 0 assert meta["raw_size"] > meta["compressed_size"] assert isinstance(meta["compressed_bytes"], bytes) class TestBinaryQuantization: """Validates FP32, FP16 half-precision, and INT8 uniform symmetric quantization.""" @pytest.fixture def engine(self) -> GradientCompressionEngine: return GradientCompressionEngine() def test_fp32_encode_decode_roundtrip(self, engine: GradientCompressionEngine) -> None: """FP32 binary serialization achieves zero error.""" weights = [1.25, -3.5, 0.0, 42.125, -0.0078125] buf = engine.encode_fp32(weights) assert len(buf) == len(weights) * 4 recovered = engine.decode_fp32(buf) assert len(recovered) == len(weights) for orig, rec in zip(weights, recovered, strict=False): assert math.isclose(orig, rec, rel_tol=1e-6) def test_fp16_encode_decode_precision(self, engine: GradientCompressionEngine) -> None: """FP16 half-precision reduces size by 50% with low reconstruction error.""" weights = [0.1234, -0.5678, 1.9876, -3.1415, 0.0] buf = engine.encode_fp16(weights) assert len(buf) == len(weights) * 2 # exactly 2 bytes per float recovered = engine.decode_fp16(buf) assert len(recovered) == len(weights) for orig, rec in zip(weights, recovered, strict=False): assert abs(orig - rec) < 1e-3 def test_int8_symmetric_quantization_bounded_error( self, engine: GradientCompressionEngine ) -> None: """INT8 symmetric affine quantization satisfies max error bound scale / 2.""" rng = np.random.default_rng(42) weights = rng.uniform(-1.0, 1.0, size=200).tolist() buf = engine.encode_int8(weights) # 4 bytes scale header + 200 bytes quantized values assert len(buf) == 4 + 200 recovered = engine.decode_int8(buf) assert len(recovered) == 200 max_val = max(abs(w) for w in weights) scale = max_val / 127.0 for orig, rec in zip(weights, recovered, strict=False): assert abs(orig - rec) <= (scale / 2.0) + 1e-5 def test_quantization_empty_and_zero_inputs(self, engine: GradientCompressionEngine) -> None: """Handles empty buffers and zero vectors gracefully.""" assert engine.encode_fp32([]) == b"" assert engine.decode_fp32(b"") == [] assert engine.encode_fp16([]) == b"" assert engine.decode_fp16(b"") == [] assert engine.encode_int8([]) == b"" assert engine.decode_int8(b"") == [] zeros = [0.0] * 10 rec_int8 = engine.decode_int8(engine.encode_int8(zeros)) assert len(rec_int8) == 10 assert all(w == 0.0 for w in rec_int8) class TestSparseCoordinateEncoding: """Validates Top-K Coordinate (COO) sparse binary format.""" @pytest.fixture def engine(self) -> GradientCompressionEngine: return GradientCompressionEngine() def test_sparse_coo_encode_decode_roundtrip(self, engine: GradientCompressionEngine) -> None: """Top-K COO buffer decodes back to dense vector with non-zeros preserved.""" weights = [0.0] * 500 weights[10] = 5.25 weights[100] = -3.50 weights[250] = 8.125 weights[499] = -1.75 buf = engine.encode_sparse_coo(weights, k_percent=0.05) # Header (10 bytes) + 4 indices (8 bytes) + 4 float16 values (8 bytes) = 26 bytes assert len(buf) < 100 recovered = engine.decode_sparse_coo(buf) assert len(recovered) == 500 assert math.isclose(recovered[10], 5.25, rel_tol=1e-3) assert math.isclose(recovered[100], -3.50, rel_tol=1e-3) assert math.isclose(recovered[250], 8.125, rel_tol=1e-3) assert math.isclose(recovered[499], -1.75, rel_tol=1e-3) # Other elements are zero assert sum(1 for w in recovered if w != 0.0) == 4 def test_sparse_coo_empty_and_invalid_buffers(self, engine: GradientCompressionEngine) -> None: """Corrupted or undersized COO buffers return safe empty lists.""" assert engine.encode_sparse_coo([]) == b"" assert engine.decode_sparse_coo(b"") == [] assert engine.decode_sparse_coo(b"CORRUPTED") == [] class TestSerializationDispatcherAndProfiling: """Validates full serialization dispatcher and wire size hierarchy.""" @pytest.fixture def engine(self) -> GradientCompressionEngine: return GradientCompressionEngine() def test_all_serialization_formats_produce_valid_roundtrip( self, engine: GradientCompressionEngine ) -> None: """All 7 SerializationFormat enums serialize and deserialize cleanly.""" rng = np.random.default_rng(42) weights = rng.normal(0.0, 0.1, size=250).tolist() for fmt in SerializationFormat: wire = engine.serialize_payload(weights, payload_format=fmt, k_percent=0.20) assert isinstance(wire, bytes) assert len(wire) > 0 rec = engine.deserialize_payload(wire, payload_format=fmt) assert len(rec) == 250 if fmt in (SerializationFormat.RAW_JSON, SerializationFormat.RAW_FP32, SerializationFormat.ZSTD_COMPRESSED): for o, r in zip(weights, rec, strict=False): assert math.isclose(o, r, abs_tol=1e-5) def test_wire_size_hierarchy(self, engine: GradientCompressionEngine) -> None: """Verifies wire size relationships: size(SPARSE_TOPK_ZSTD) < size(INT8) < size(FP16) < size(FP32) < size(RAW_JSON). """ rng = np.random.default_rng(42) weights = rng.normal(0.0, 0.05, size=1969).tolist() profiles = engine.compare_all_serialization_formats(weights, k_percent=0.20) p_json = profiles[SerializationFormat.RAW_JSON.value] p_fp32 = profiles[SerializationFormat.RAW_FP32.value] p_fp16 = profiles[SerializationFormat.QUANTIZED_FP16.value] p_int8 = profiles[SerializationFormat.QUANTIZED_INT8.value] p_sparse_zstd = profiles[SerializationFormat.SPARSE_TOPK_ZSTD.value] assert p_json.wire_bytes > p_fp32.wire_bytes assert p_fp32.wire_bytes > p_fp16.wire_bytes assert p_fp16.wire_bytes > p_int8.wire_bytes assert p_int8.wire_bytes > p_sparse_zstd.wire_bytes # FP32 on 1,969 parameters is exactly 7,876 bytes assert p_fp32.wire_bytes == 1969 * 4 # FP16 is exactly half of FP32 assert p_fp16.wire_bytes == 1969 * 2 # INT8 is 1 byte per param + 4 bytes header assert p_int8.wire_bytes == 1969 + 4 def test_profile_serialized_payload_metrics(self, engine: GradientCompressionEngine) -> None: """SerializedPayloadProfile fields are consistent and bounded.""" weights = [float(i) for i in range(100)] profile = engine.profile_serialized_payload( weights, payload_format=SerializationFormat.QUANTIZED_FP16 ) assert profile.num_parameters == 100 assert profile.raw_uncompressed_bytes == 400 assert profile.wire_bytes == 200 assert profile.compression_ratio_vs_fp32 == 0.50 assert profile.bandwidth_savings_pct == 50.0 assert profile.serialization_time_us >= 0.0 assert profile.deserialization_time_us >= 0.0 assert profile.reconstruction_mae >= 0.0 class TestFederatedCommunicationProfiler: """Validates multi-round communication accounting across federated aggregation protocols.""" @pytest.fixture def profiler(self) -> FederatedCommunicationProfiler: return FederatedCommunicationProfiler() def test_fedavg_and_fedprox_identical_bandwidth( self, profiler: FederatedCommunicationProfiler ) -> None: """FedProx has identical communication payload to FedAvg.""" p_avg = profiler.profile_protocol( num_parameters=1969, num_clients=3, num_rounds=5, protocol=FederatedProtocol.FED_AVG ) p_prox = profiler.profile_protocol( num_parameters=1969, num_clients=3, num_rounds=5, protocol=FederatedProtocol.FED_PROX ) assert p_avg.total_round_bytes == p_prox.total_round_bytes assert p_avg.total_multi_round_bytes == p_prox.total_multi_round_bytes assert p_avg.relative_overhead_vs_fedavg == 1.0 assert p_prox.relative_overhead_vs_fedavg == 1.0 # 1 server broadcast (7,876) + 3 client uploads (3 * 7,876) = 31,504 bytes per round assert p_avg.total_round_bytes == 31504 assert p_avg.total_multi_round_bytes == 31504 * 5 def test_scaffold_exact_2x_overhead(self, profiler: FederatedCommunicationProfiler) -> None: """SCAFFOLD requires exactly 2x bandwidth due to control variates.""" p_avg = profiler.profile_protocol( num_parameters=1969, num_clients=3, num_rounds=5, protocol=FederatedProtocol.FED_AVG ) p_scaffold = profiler.profile_protocol( num_parameters=1969, num_clients=3, num_rounds=5, protocol=FederatedProtocol.SCAFFOLD ) assert p_scaffold.total_round_bytes == 2 * p_avg.total_round_bytes assert p_scaffold.total_multi_round_bytes == 2 * p_avg.total_multi_round_bytes assert p_scaffold.relative_overhead_vs_fedavg == 2.0 def test_curve25519_secagg_overhead_bounds( self, profiler: FederatedCommunicationProfiler ) -> None: """Curve25519 SecAgg adds minor cryptographic coordination overhead (4% - 10%).""" p_secagg = profiler.profile_protocol( num_parameters=1969, num_clients=3, num_rounds=5, protocol=FederatedProtocol.CURVE25519_SECAGG, ) assert p_secagg.relative_overhead_vs_fedavg > 1.0 assert p_secagg.relative_overhead_vs_fedavg < 1.15 assert p_secagg.total_multi_round_bytes > 157520 def test_tenseal_ckks_expansion_bounds( self, profiler: FederatedCommunicationProfiler ) -> None: """TenSEAL CKKS exhibits polynomial ciphertext expansion (~10.5x).""" p_ckks = profiler.profile_protocol( num_parameters=1969, num_clients=3, num_rounds=5, protocol=FederatedProtocol.TENSEAL_CKKS, ) assert p_ckks.relative_overhead_vs_fedavg == 10.5 assert p_ckks.total_multi_round_mb > 1.0 def test_run_full_bandwidth_benchmark_suite( self, profiler: FederatedCommunicationProfiler ) -> None: """Executes complete suite across all 7 formats and 5 protocols.""" suite = profiler.run_full_bandwidth_benchmark( num_clients=3, num_rounds=5, model_name="CanonicalFraudMLP" ) assert isinstance(suite, BandwidthBenchmarkSuiteResult) assert suite.num_parameters == 1969 assert len(suite.serialization_profiles) == len(SerializationFormat) assert len(suite.protocol_profiles) == len(FederatedProtocol) assert "RAW_FP32" in suite.serialization_profiles assert "FED_AVG" in suite.protocol_profiles class TestFederatedEngineIntegration: """Validates live FederatedLearningEngine communication telemetry integration.""" def test_fl_engine_profile_communication_round(self) -> None: """FederatedLearningEngine delegates communication profiling to profiler.""" mock_settings = MagicMock() mock_model_svc = MagicMock() mock_privacy_svc = MagicMock() engine = FederatedLearningEngine( settings=mock_settings, model_service=mock_model_svc, privacy_service=mock_privacy_svc, ) # Profile FedAvg round profile_avg = engine.profile_communication_round( num_parameters=1969, num_clients=3, method=AggregationMethod.FED_AVG, serialization_format=SerializationFormat.RAW_FP32, ) assert isinstance(profile_avg, ProtocolCommunicationProfile) assert profile_avg.protocol == FederatedProtocol.FED_AVG assert profile_avg.total_round_bytes == 31504 # Profile SCAFFOLD round profile_scaffold = engine.profile_communication_round( num_parameters=1969, num_clients=3, method=AggregationMethod.SCAFFOLD, serialization_format=SerializationFormat.RAW_FP32, ) assert profile_scaffold.protocol == FederatedProtocol.SCAFFOLD assert profile_scaffold.total_round_bytes == 63008