test: add unit tests for columns

src/columns.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, make_dataclass
 
 
-def fields(raw_class):
+def _fields(raw_class):
     return [v for k, v in raw_class.__dict__.items() if k[:2] != "__" and k[-2:] != "__"]
 
 
@@ -19,45 +19,104 @@ class ColumnContent:
 
 def get_default_auto_eval_column_dict():
     auto_eval_column_dict = []
-    # Init
-    auto_eval_column_dict.append(["rank", ColumnContent, ColumnContent(COL_NAME_RANK, "number", True)])
+    auto_eval_column_dict.append(
+        [
+            "rank",
+            ColumnContent,
+            ColumnContent(
+                COL_NAME_RANK,
+                "number",
+                True
+            )
+        ]
+    )
     auto_eval_column_dict.append(
         [
             "retrieval_model",
             ColumnContent,
-            ColumnContent(COL_NAME_RETRIEVAL_MODEL, "markdown", True, hidden=False, never_hidden=True),
+            ColumnContent(
+                COL_NAME_RETRIEVAL_MODEL,
+                "markdown",
+                True,
+                never_hidden=True
+            )
         ]
     )
     auto_eval_column_dict.append(
         [
             "reranking_model",
             ColumnContent,
-            ColumnContent(COL_NAME_RERANKING_MODEL, "markdown", True, hidden=False, never_hidden=True),
+            ColumnContent(
+                COL_NAME_RERANKING_MODEL,
+                "markdown",
+                True,
+                never_hidden=True
+            )
         ]
     )
     auto_eval_column_dict.append(
-        ["revision", ColumnContent, ColumnContent(COL_NAME_REVISION, "markdown", True, never_hidden=True)]
+        [
+            "revision",
+            ColumnContent,
+            ColumnContent(
+                COL_NAME_REVISION,
+                "markdown",
+                True,
+                never_hidden=True
+            )
+        ]
     )
     auto_eval_column_dict.append(
-        ["timestamp", ColumnContent, ColumnContent(COL_NAME_TIMESTAMP, "date", True, never_hidden=True)]
+        [
+            "timestamp",
+            ColumnContent,
+            ColumnContent(
+                COL_NAME_TIMESTAMP, "date", True, never_hidden=True
+            )
+        ]
+    )
+    auto_eval_column_dict.append(
+        [
+            "average",
+            ColumnContent,
+            ColumnContent(COL_NAME_AVG, "number", True)
+        ]
     )
-    auto_eval_column_dict.append(["average", ColumnContent, ColumnContent(COL_NAME_AVG, "number", True)])
     auto_eval_column_dict.append(
         [
             "retrieval_model_link",
             ColumnContent,
-            ColumnContent(COL_NAME_RETRIEVAL_MODEL_LINK, "markdown", False, hidden=True, never_hidden=False),
+            ColumnContent(
+                COL_NAME_RETRIEVAL_MODEL_LINK,
+                "markdown",
+                False,
+                hidden=True,
+            )
         ]
     )
     auto_eval_column_dict.append(
         [
             "reranking_model_link",
             ColumnContent,
-            ColumnContent(COL_NAME_RERANKING_MODEL_LINK, "markdown", False, hidden=True, never_hidden=False),
+            ColumnContent(
+                COL_NAME_RERANKING_MODEL_LINK,
+                "markdown",
+                False,
+                hidden=True,
+            )
         ]
     )
     auto_eval_column_dict.append(
-        ["is_anonymous", ColumnContent, ColumnContent(COL_NAME_IS_ANONYMOUS, "bool", False, hidden=True)]
+        [
+            "is_anonymous",
+            ColumnContent,
+            ColumnContent(
+                COL_NAME_IS_ANONYMOUS,
+                "bool",
+                False,
+                hidden=True
+            )
+        ]
     )
     return auto_eval_column_dict
 
@@ -76,8 +135,8 @@ def make_autoevalcolumn(cls_name, benchmarks):
 
 def get_default_col_names_and_types(benchmarks):
     AutoEvalColumn = make_autoevalcolumn("AutoEvalColumn", benchmarks)
-    col_names = [c.name for c in fields(AutoEvalColumn) if not c.hidden]
-    col_types = [c.type for c in fields(AutoEvalColumn) if not c.hidden]
+    col_names = [c.name for c in _fields(AutoEvalColumn) if not c.hidden]
+    col_types = [c.type for c in _fields(AutoEvalColumn) if not c.hidden]
     return col_names, col_types
 
 

tests/src/test_columns.py

@@ -0,0 +1,116 @@
+import pytest
+
+from src.benchmarks import QABenchmarks, LongDocBenchmarks
+from src.columns import get_default_auto_eval_column_dict, \
+    get_fixed_col_names_and_types, get_default_col_names_and_types, make_autoevalcolumn, COL_NAME_RANK, \
+    COL_NAME_RETRIEVAL_MODEL, \
+    COL_NAME_RERANKING_MODEL, COL_NAME_REVISION, \
+    COL_NAME_TIMESTAMP, COL_NAME_AVG
+
+
+# Ref: https://github.com/AIR-Bench/AIR-Bench/blob/4b27b8a8f2047a963805fcf6fb9d74be51ec440c/docs/available_tasks.md
+# 24.05
+# | Task | dev | test |
+# | ---- | --- | ---- |
+# | Long-Doc | 4 | 11 |
+# | QA | 54 | 53 |
+#
+# 24.04
+# | Task | test |
+# | ---- | ---- |
+# | Long-Doc | 15 |
+# | QA | 13 |
+
+@pytest.fixture()
+def expected_col_names():
+    return [
+        "rank",
+        "retrieval_model",
+        "reranking_model",
+        "revision",
+        "timestamp",
+        "average",
+        "retrieval_model_link",
+        "reranking_model_link",
+        "is_anonymous",
+    ]
+
+
+@pytest.fixture()
+def expected_hidden_col_names():
+    return [
+        "retrieval_model_link",
+        "reranking_model_link",
+        "is_anonymous",
+    ]
+
+
+def test_get_default_auto_eval_column_dict(
+        expected_col_names, expected_hidden_col_names):
+    col_list = get_default_auto_eval_column_dict()
+    assert len(col_list) == 9
+    hidden_cols = []
+    for col_tuple, expected_col in zip(col_list, expected_col_names):
+        col, _, col_content = col_tuple
+        assert col == expected_col
+        if col_content.hidden:
+            hidden_cols.append(col)
+    assert hidden_cols == expected_hidden_col_names
+
+
+def test_get_fixed_col_names_and_types():
+    col_names, col_types = get_fixed_col_names_and_types()
+    assert len(col_names) == 6
+    assert len(col_types) == 6
+    expected_col_and_type = [
+        (COL_NAME_RANK, "number"),
+        (COL_NAME_RETRIEVAL_MODEL, "markdown"),
+        (COL_NAME_RERANKING_MODEL, "markdown"),
+        (COL_NAME_REVISION, "markdown"),
+        (COL_NAME_TIMESTAMP, "date"),
+        (COL_NAME_AVG, "number"),
+    ]
+    for col_name, col_type, (c_name, c_type) in zip(col_names, col_types, expected_col_and_type):
+        assert col_name == c_name
+        assert col_type == c_type
+
+
+@pytest.mark.parametrize(
+    'benchmarks, expected_benchmark_len',
+    [
+        (QABenchmarks, {"air_bench_2404": 13, "air_bench_2405": 53}),
+        (LongDocBenchmarks, {"air_bench_2404": 15, "air_bench_2405": 11})
+    ]
+)
+def test_make_autoevalcolumn(
+        benchmarks, expected_benchmark_len, expected_col_names):
+    expected_default_attrs = frozenset(expected_col_names)
+    for benchmark in benchmarks:
+        TestEvalColumn = make_autoevalcolumn("TestEvalColumn", benchmark)
+        attrs = []
+        for k, v in TestEvalColumn.__dict__.items():
+            if not k.startswith("__"):
+                attrs.append(k)
+        attrs = frozenset(attrs)
+        assert expected_default_attrs.issubset(attrs)
+        benchmark_attrs = attrs.difference(expected_default_attrs)
+        assert len(benchmark_attrs) == expected_benchmark_len[benchmark.name]
+
+
+@pytest.mark.parametrize(
+    'benchmarks, expected_benchmark_len',
+    [
+        (QABenchmarks, {"air_bench_2404": 13, "air_bench_2405": 53}),
+        (LongDocBenchmarks, {"air_bench_2404": 15, "air_bench_2405": 11})
+    ]
+)
+def test_get_default_col_names_and_types(
+        benchmarks,
+        expected_benchmark_len,
+        expected_col_names,
+        expected_hidden_col_names):
+    default_col_len = len(expected_col_names)
+    hidden_col_len = len(expected_hidden_col_names)
+    for benchmark in benchmarks:
+        col_names, col_types = get_default_col_names_and_types(benchmark)
+        assert len(col_names) == expected_benchmark_len[benchmark.name] + default_col_len - hidden_col_len