Add retry_mechanism_analysis/tests/test_core.py

retry_mechanism_analysis/tests/test_core.py

@@ -0,0 +1,72 @@
+import pytest
+import pandas as pd
+from retry_mechanism_analysis import core
+
+
+def _sample_log_data():
+    # Example dataset with two parallelism levels
+    return [
+        {'parallelism_level': 2, 'retry_latency': 100, 'stratum': 'A'},
+        {'parallelism_level': 2, 'retry_latency': 200, 'stratum': 'A'},
+        {'parallelism_level': 2, 'retry_latency': 400, 'stratum': 'B'},
+        {'parallelism_level': 4, 'retry_latency': 150, 'stratum': 'A'},
+        {'parallelism_level': 4, 'retry_latency': 250, 'stratum': 'A'},
+        {'parallelism_level': 4, 'retry_latency': 500, 'stratum': 'B'}
+    ]
+
+
+def _expected_metrics(level, latencies):
+    series = pd.Series(sorted(latencies))
+    p50 = series.quantile(0.5)
+    p95 = series.quantile(0.95)
+    p99 = series.quantile(0.99)
+    return {
+        'parallelism_level': level,
+        'p50': pytest.approx(p50, rel=1e-3),
+        'p95': pytest.approx(p95, rel=1e-3),
+        'p99': pytest.approx(p99, rel=1e-3)
+    }
+
+
+def test_compare_retry_overhead_basics():
+    log_data = _sample_log_data()
+    result = core.compare_retry_overhead(log_data)
+    
+    assert isinstance(result, list)
+    assert all(set(r.keys()) == {'parallelism_level', 'p50', 'p95', 'p99'} for r in result)
+    
+    level2 = next(r for r in result if r['parallelism_level'] == 2)
+    level4 = next(r for r in result if r['parallelism_level'] == 4)
+
+    lat2 = [100, 200, 400]
+    lat4 = [150, 250, 500]
+
+    expected2 = _expected_metrics(2, lat2)
+    expected4 = _expected_metrics(4, lat4)
+
+    for key in ('p50', 'p95', 'p99'):
+        assert level2[key] == expected2[key]
+        assert level4[key] == expected4[key]
+
+
+def test_compare_retry_overhead_invalid_input():
+    with pytest.raises((TypeError, ValueError)):
+        core.compare_retry_overhead(None)
+
+    # Invalid type in list
+    bad_data = [{'parallelism_level': 'x', 'retry_latency': 'notnum'}]
+    with pytest.raises((TypeError, ValueError)):
+        core.compare_retry_overhead(bad_data)
+
+
+def test_compare_retry_overhead_empty_list():
+    result = core.compare_retry_overhead([])
+    assert result == []
+
+
+def test_compare_retry_overhead_dataframe_support():
+    # Test if the implementation can handle pandas DataFrame input (optional convenience)
+    df = pd.DataFrame(_sample_log_data())
+    result = core.compare_retry_overhead(df.to_dict(orient='records'))
+    assert isinstance(result, list)
+    assert all('p50' in r for r in result)