test: add unit tests for loaders

src/loaders.py

@@ -53,7 +53,7 @@ def load_raw_eval_results(results_path: Union[Path, str]) -> List[FullEvalResult
 
 
 def load_leaderboard_datastore(file_path, version) -> LeaderboardDataStore:
-    ds = LeaderboardDataStore(version, get_safe_name(version), None, None, None, None, None, None, None, None)
+    ds = LeaderboardDataStore(version, get_safe_name(version))
     ds.raw_data = load_raw_eval_results(file_path)
     print(f"raw data: {len(ds.raw_data)}")
 
@@ -61,7 +61,9 @@ def load_leaderboard_datastore(file_path, version) -> LeaderboardDataStore:
     print(f"QA data loaded: {ds.qa_raw_df.shape}")
     ds.qa_fmt_df = ds.qa_raw_df.copy()
     qa_cols, ds.qa_types = get_default_cols(TaskType.qa, ds.slug, add_fix_cols=True)
+    # by default, drop the anonymous submissions
     ds.qa_fmt_df = ds.qa_fmt_df[~ds.qa_fmt_df[COL_NAME_IS_ANONYMOUS]][qa_cols]
+    # reset the rank after dropping the anonymous submissions
     ds.qa_fmt_df = reset_rank(ds.qa_fmt_df)
     ds.qa_fmt_df.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
 
@@ -69,7 +71,9 @@ def load_leaderboard_datastore(file_path, version) -> LeaderboardDataStore:
     print(f"Long-Doc data loaded: {len(ds.doc_raw_df)}")
     ds.doc_fmt_df = ds.doc_raw_df.copy()
     doc_cols, ds.doc_types = get_default_cols(TaskType.long_doc, ds.slug, add_fix_cols=True)
+    # by default, drop the anonymous submissions
     ds.doc_fmt_df = ds.doc_fmt_df[~ds.doc_fmt_df[COL_NAME_IS_ANONYMOUS]][doc_cols]
+    # reset the rank after dropping the anonymous submissions
     ds.doc_fmt_df = reset_rank(ds.doc_fmt_df)
     ds.doc_fmt_df.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
 

tests/src/test_loaders.py

@@ -0,0 +1,54 @@
+import pandas as pd
+import pytest
+from pathlib import Path
+
+from src.loaders import load_raw_eval_results, load_leaderboard_datastore, load_eval_results
+
+cur_fp = Path(__file__)
+
+
+@pytest.mark.parametrize(
+    "version",
+    ["AIR-Bench_24.04", "AIR-Bench_24.05"]
+)
+def test_load_raw_eval_results(version):
+    raw_data = load_raw_eval_results(
+        cur_fp.parents[1] / f"toydata/eval_results/{version}"
+    )
+    assert len(raw_data) == 1
+    full_eval_result = raw_data[0]
+    expected_attr = [
+        'eval_name',
+        'retrieval_model',
+        'reranking_model',
+        'retrieval_model_link',
+        'reranking_model_link',
+        'results',
+        'timestamp',
+        'revision',
+        'is_anonymous'
+    ]
+    result_attr = [k for k in full_eval_result.__dict__.keys() if k[:2] != "__" and k[-2:] != "__"]
+    assert sorted(expected_attr) == sorted(result_attr)
+
+
+@pytest.mark.parametrize(
+    "version",
+    ["AIR-Bench_24.04", "AIR-Bench_24.05"]
+)
+def test_load_leaderboard_datastore(version):
+    file_path = cur_fp.parents[1] / f"toydata/eval_results/{version}"
+    datastore = load_leaderboard_datastore(file_path, version)
+    for k, v in datastore.__dict__.items():
+        if k[:2] != "__" and k[-2:] != "__":
+            if isinstance(v, list):
+                assert v
+            elif isinstance(v, pd.DataFrame):
+                assert not v.empty
+
+
+def test_load_eval_results():
+    file_path = cur_fp.parents[1] / "toydata/eval_results/"
+    datastore_dict = load_eval_results(file_path)
+    assert len(datastore_dict) == 2
+