update unit tests to cover milestone data

Author: ngc92

unit-tests/test_leaderboard_db.py

@@ -1,8 +1,11 @@
 import copy
 import dataclasses
 import datetime
+import decimal
 import subprocess
 import time
+from pathlib import Path
+from unittest.mock import ANY
 
 import pytest
 from test_report import sample_compile_result, sample_run_result, sample_system_info
@@ -16,8 +19,14 @@
 
 @pytest.fixture(scope="module")
 def docker_compose():
+    tgt_path = Path.cwd()
+    if tgt_path.name == "unit-tests":
+        tgt_path = tgt_path.parent
+
     """Start a test database and run migrations"""
-    subprocess.check_call(["docker", "compose", "-f", "docker-compose.test.yml", "up", "-d"])
+    subprocess.check_call(
+        ["docker", "compose", "-f", "docker-compose.test.yml", "up", "-d"], cwd=tgt_path
+    )
 
     try:
         # Wait for migrations to finish
@@ -26,6 +35,7 @@ def docker_compose():
                 ["docker", "compose", "-f", "docker-compose.test.yml", "ps", "-q", "migrate-test"],
                 capture_output=True,
                 text=True,
+                cwd=tgt_path,
             )
 
             if not result.stdout.strip():  # Container no longer exists
@@ -37,6 +47,7 @@ def docker_compose():
             ["docker", "compose", "-f", "docker-compose.test.yml", "logs", "migrate-test"],
             capture_output=True,
             text=True,
+            cwd=tgt_path,
         )
 
         if "error" in logs.stdout.lower():
@@ -52,7 +63,9 @@ def docker_compose():
             ssl_mode="disable",
         )
     finally:
-        subprocess.run(["docker", "compose", "-f", "docker-compose.test.yml", "down", "-v"])
+        subprocess.run(
+            ["docker", "compose", "-f", "docker-compose.test.yml", "down", "-v"], cwd=tgt_path
+        )
 
 
 def _nuke_contents(db):
@@ -114,7 +127,7 @@ def _create_submission_run(
 ):
     """Creates a submission run with suitable default values"""
     db.create_submission_run(
-        submission,
+        submission=submission,
         start=start or datetime.datetime.now(tz=datetime.timezone.utc),
         end=end
         or (datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=10)),
@@ -268,9 +281,9 @@ def test_leaderboard_submission_basic(database, submit_leaderboard):
     with database as db:
         end_time = submit_time + datetime.timedelta(seconds=10)
         db.create_submission_run(
-            sub_id,
-            submit_time,
-            end_time,
+            submission=sub_id,
+            start=submit_time,
+            end=end_time,
             mode="test",
             secret=False,
             runner="A100",
@@ -282,9 +295,9 @@ def test_leaderboard_submission_basic(database, submit_leaderboard):
         # run ends after the contest deadline; this is valid
         end_time_2 = submit_time + datetime.timedelta(days=1, hours=1)
         db.create_submission_run(
-            sub_id,
-            submit_time,
-            end_time_2,
+            submission=sub_id,
+            start=submit_time,
+            end=end_time_2,
             mode="leaderboard",
             secret=True,
             runner="H100",
@@ -577,6 +590,107 @@ def test_leaderboard_update(database, task_directory):
         }
 
 
+def test_leaderboard_milestones(database, submit_leaderboard):
+    with database as db:
+        lb_id = db.get_leaderboard_id("submit-leaderboard")
+        milestones = db.get_leaderboard_milestones(lb_id)
+        assert milestones == []
+
+        # at this point, created_at is filled in at the DB level,
+        # so we cannot set a fixed value for it in the tests below
+        db.create_milestone(lb_id, "Milestone", "sample code", "Test milestone")
+        db.create_milestone(
+            lb_id, "Milestone2", "other code", "Second milestone", exclude_gpus=["T4"]
+        )
+        milestones = db.get_leaderboard_milestones(lb_id)
+        assert milestones == [
+            {
+                "code": "sample code",
+                "created_at": ANY,
+                "description": "Test milestone",
+                "exclude_gpus": [""],
+                "id": 1,
+                "name": "Milestone",
+            },
+            {
+                "code": "other code",
+                "created_at": ANY,
+                "description": "Second milestone",
+                "exclude_gpus": ["T4"],
+                "id": 2,
+                "name": "Milestone2",
+            },
+        ]
+
+        db.delete_milestones(lb_id)
+        milestones = db.get_leaderboard_milestones(lb_id)
+        assert milestones == []
+
+
+def test_leaderboard_milestone_runs(database, submit_leaderboard):
+    with database as db:
+        lb_id = db.get_leaderboard_id("submit-leaderboard")
+        ms_id = db.create_milestone(lb_id, "Milestone", "sample code", "Test milestone")
+
+        start = datetime.datetime.now(tz=datetime.timezone.utc)
+        end = start + datetime.timedelta(seconds=10)
+        db.create_submission_run(
+            milestone=ms_id,
+            start=start,
+            end=end,
+            mode="leaderboard",
+            secret=False,
+            runner="A100",
+            score=5,
+            compilation=None,
+            result=sample_run_result(),
+            system=sample_system_info(),
+        )
+
+        runs = db.get_runs_generic(milestone_id=ms_id)
+        assert runs == [
+            {
+                "compilation": None,
+                "start_time": start,
+                "end_time": end,
+                "meta": {
+                    "command": "./test",
+                    "duration": 1.5,
+                    "exit_code": 0,
+                    "stderr": "",
+                    "stdout": "All tests passed",
+                    "success": True,
+                },
+                "mode": "leaderboard",
+                "passed": True,
+                "result": {
+                    "test-count": "3",
+                    "test.0.message": "Addition works correctly",
+                    "test.0.spec": "Test addition",
+                    "test.0.status": "pass",
+                    "test.1.spec": "Test multiplication",
+                    "test.1.status": "pass",
+                    "test.2.error": "Division by zero",
+                    "test.2.spec": "Test division",
+                    "test.2.status": "fail",
+                },
+                "runner": "A100",
+                "score": decimal.Decimal("5"),
+                "secret": False,
+                "system": {
+                    "cpu": "Intel i9-12900K",
+                    "gpu": "NVIDIA RTX 4090",
+                    "platform": "Linux-5.15.0",
+                    "torch": "2.0.1+cu118",
+                },
+            }
+        ]
+
+        db.delete_milestone_runs(lb_id)
+        runs = db.get_runs_generic(milestone_id=ms_id)
+        assert runs == []
+
+
 def test_generate_stats(database, submit_leaderboard):
     with database as db:
         start = datetime.datetime.now(tz=datetime.timezone.utc)

unit-tests/test_task.py

@@ -10,6 +10,7 @@
     Language,
     LeaderboardDefinition,
     LeaderboardTask,
+    MilestoneData,
     PythonTaskData,
     RankCriterion,
     build_task_config,
@@ -57,7 +58,6 @@ def test_from_dict_python_task():
 
 
 def test_from_dict_cuda_task():
-    """Test creating LeaderboardTask from dict with CUDA config"""
     """Test creating LeaderboardTask from dict with CUDA config"""
     data = {
         "lang": "cu",
@@ -93,7 +93,7 @@ def test_type_mismatch():
         )
 
 
-def test_to_dict(leaderboard_task):
+def test_to_dict(leaderboard_task: LeaderboardTask):
     """Test converting LeaderboardTask to dict"""
     result = leaderboard_task.to_dict()
 
@@ -114,15 +114,15 @@ def test_to_dict(leaderboard_task):
     ]
 
 
-def test_serialization_roundtrip(leaderboard_task):
+def test_serialization_roundtrip(leaderboard_task: LeaderboardTask):
     """Test to_str and from_str work together"""
     json_str = leaderboard_task.to_str()
     reconstructed = LeaderboardTask.from_str(json_str)
 
     assert reconstructed == leaderboard_task
 
 
-def test_build_task_config_python(leaderboard_task):
+def test_build_task_config_python(leaderboard_task: LeaderboardTask):
     """Test build_task_config with Python task and submission content."""
     submission_content = "print('Hello World')"
     arch = "sm_80"
@@ -235,6 +235,11 @@ def test_build_task_config_cuda():
 templates:
   Python: "template.py"
   CUDA: "template.cu"
+milestones:
+  - name: "Milestone"
+    source: "milestone.py"
+    description: "This milestone is a test milestone"
+    exclude_gpus: ["A100"]
 """
 
 
@@ -245,13 +250,14 @@ def task_directory(tmp_path):
     Path.write_text(tmp_path / "kernel.py", "def kernel(): pass")
     Path.write_text(tmp_path / "template.py", "# Python template")
     Path.write_text(tmp_path / "template.cu", "// CUDA template")
+    Path.write_text(tmp_path / "milestone.py", "def milestone(): pass")
 
     # Create task.yml
     Path.write_text(tmp_path / "task.yml", TASK_YAML)
     return tmp_path
 
 
-def test_make_task_definition(task_directory):
+def test_make_task_definition(task_directory: Path):
     """Test make_task_definition with a complete YAML structure"""
 
     # Test the function
@@ -261,6 +267,14 @@ def test_make_task_definition(task_directory):
     assert isinstance(result, LeaderboardDefinition)
     assert result.description == "Test task description"
     assert result.templates == {"Python": "# Python template", "CUDA": "// CUDA template"}
+    assert result.milestones == [
+        MilestoneData(
+            name="Milestone",
+            code="def milestone(): pass",
+            description="This milestone is a test milestone",
+            exclude_gpus=["A100"],
+        )
+    ]
 
     # Verify the task
     task = result.task