update unit tests to cover milestone data

Author: ngc92

unit-tests/test_task.py

@@ -13,7 +13,7 @@
     PythonTaskData,
     RankCriterion,
     build_task_config,
-    make_task_definition,
+    make_task_definition, MilestoneData,
 )
 
 
@@ -57,7 +57,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 +92,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 +113,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"
@@ -180,6 +179,11 @@ def test_build_task_config_python(leaderboard_task):
 templates:
   Python: "template.py"
   CUDA: "template.cu"
+milestones:
+  - name: "Milestone"
+    source: "milestone.py"
+    description: "This milestone is a test milestone"
+    exclude_gpus: ["A100"]
 """
 
 
@@ -190,13 +194,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
@@ -206,6 +211,7 @@ 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