diff --git a/src/xer_mcp/db/queries.py b/src/xer_mcp/db/queries.py index 18b6297..4b47d32 100644 --- a/src/xer_mcp/db/queries.py +++ b/src/xer_mcp/db/queries.py @@ -396,9 +396,9 @@ def get_project_summary(project_id: str) -> dict | None: cur.execute("SELECT COUNT(*) FROM activities") activity_count = cur.fetchone()[0] - # Get milestone count + # Get milestone count (both start and finish milestones) with db.cursor() as cur: - cur.execute("SELECT COUNT(*) FROM activities WHERE task_type = 'TT_Mile'") + cur.execute("SELECT COUNT(*) FROM activities WHERE task_type IN ('TT_Mile', 'TT_FinMile')") milestone_count = cur.fetchone()[0] # Get critical activity count @@ -419,16 +419,16 @@ def get_project_summary(project_id: str) -> dict | None: def query_milestones() -> list[dict]: - """Query all milestone activities. + """Query all milestone activities (both start and finish milestones). Returns: - List of milestone activity dicts + List of milestone activity dicts with milestone_type (start/finish) """ query = """ SELECT task_id, task_code, task_name, - target_start_date, target_end_date, status_code + target_start_date, target_end_date, status_code, milestone_type FROM activities - WHERE task_type = 'TT_Mile' + WHERE task_type IN ('TT_Mile', 'TT_FinMile') ORDER BY target_start_date, task_code """ @@ -444,6 +444,7 @@ def query_milestones() -> list[dict]: "target_start_date": row[3], "target_end_date": row[4], "status_code": row[5], + "milestone_type": row[6], } for row in rows ] diff --git a/src/xer_mcp/parser/table_handlers/task.py b/src/xer_mcp/parser/table_handlers/task.py index 631c66a..962fd82 100644 --- a/src/xer_mcp/parser/table_handlers/task.py +++ b/src/xer_mcp/parser/table_handlers/task.py @@ -26,6 +26,22 @@ class TaskHandler(TableHandler): float_str = data.get("total_float_hr_cnt", "") total_float = float(float_str) if float_str else None + # Parse milestone_type + # First check explicit milestone_type field (MS_Start -> 'start', MS_Finish -> 'finish') + # Then derive from task_type (TT_Mile -> 'start', TT_FinMile -> 'finish') + raw_milestone_type = data.get("milestone_type", "") + task_type = data.get("task_type", "") + milestone_type = None + if raw_milestone_type: + if raw_milestone_type == "MS_Start": + milestone_type = "start" + elif raw_milestone_type == "MS_Finish": + milestone_type = "finish" + elif task_type == "TT_Mile": + milestone_type = "start" + elif task_type == "TT_FinMile": + milestone_type = "finish" + return { "task_id": data.get("task_id", ""), "proj_id": data.get("proj_id", ""), @@ -42,4 +58,5 @@ class TaskHandler(TableHandler): "act_end_date": convert_date(data.get("act_end_date")), "total_float_hr_cnt": total_float, "driving_path_flag": driving_path, + "milestone_type": milestone_type, } diff --git a/tests/unit/test_table_handlers.py b/tests/unit/test_table_handlers.py index ea0b3f2..838746c 100644 --- a/tests/unit/test_table_handlers.py +++ b/tests/unit/test_table_handlers.py @@ -157,6 +157,147 @@ class TestTaskHandler: assert result["early_start_date"] is None assert result["early_end_date"] is None + def test_parse_milestone_type_start(self) -> None: + """Handler should parse milestone_type='start' for start milestones.""" + from xer_mcp.parser.table_handlers.task import TaskHandler + + handler = TaskHandler() + + fields = [ + "task_id", + "proj_id", + "task_code", + "task_name", + "task_type", + "milestone_type", + ] + values = [ + "2001", + "1001", + "MS-START", + "Project Start", + "TT_Mile", + "MS_Start", + ] + + result = handler.parse_row(fields, values) + + assert result is not None + assert result["task_type"] == "TT_Mile" + assert result["milestone_type"] == "start" + + def test_parse_milestone_type_finish(self) -> None: + """Handler should parse milestone_type='finish' for finish milestones.""" + from xer_mcp.parser.table_handlers.task import TaskHandler + + handler = TaskHandler() + + fields = [ + "task_id", + "proj_id", + "task_code", + "task_name", + "task_type", + "milestone_type", + ] + values = [ + "2002", + "1001", + "MS-END", + "Project Complete", + "TT_Mile", + "MS_Finish", + ] + + result = handler.parse_row(fields, values) + + assert result is not None + assert result["task_type"] == "TT_Mile" + assert result["milestone_type"] == "finish" + + def test_parse_milestone_type_null_for_non_milestones(self) -> None: + """Handler should return None for milestone_type on non-milestone activities.""" + from xer_mcp.parser.table_handlers.task import TaskHandler + + handler = TaskHandler() + + fields = [ + "task_id", + "proj_id", + "task_code", + "task_name", + "task_type", + ] + values = [ + "2003", + "1001", + "A1000", + "Regular Task", + "TT_Task", + ] + + result = handler.parse_row(fields, values) + + assert result is not None + assert result["task_type"] == "TT_Task" + assert result["milestone_type"] is None + + def test_parse_milestone_type_derived_from_tt_mile(self) -> None: + """Handler should derive milestone_type='start' from task_type TT_Mile.""" + from xer_mcp.parser.table_handlers.task import TaskHandler + + handler = TaskHandler() + + # No explicit milestone_type field - derive from task_type + fields = [ + "task_id", + "proj_id", + "task_code", + "task_name", + "task_type", + ] + values = [ + "2001", + "1001", + "MS-START", + "Project Start", + "TT_Mile", + ] + + result = handler.parse_row(fields, values) + + assert result is not None + assert result["task_type"] == "TT_Mile" + assert result["milestone_type"] == "start" + + def test_parse_milestone_type_derived_from_tt_finmile(self) -> None: + """Handler should derive milestone_type='finish' from task_type TT_FinMile.""" + from xer_mcp.parser.table_handlers.task import TaskHandler + + handler = TaskHandler() + + # TT_FinMile is a finish milestone + fields = [ + "task_id", + "proj_id", + "task_code", + "task_name", + "task_type", + ] + values = [ + "2002", + "1001", + "MS-END", + "Project Complete", + "TT_FinMile", + ] + + result = handler.parse_row(fields, values) + + assert result is not None + assert result["task_type"] == "TT_FinMile" + assert result["milestone_type"] == "finish" + class TestTaskpredHandler: """Tests for TASKPRED table handler."""