diff --git a/models.py b/models.py index ebb279b..46d92f6 100644 --- a/models.py +++ b/models.py @@ -1,5 +1,6 @@ from dataclasses import dataclass, field -from typing import List, Optional +from typing import List, Optional, Dict, Any +from datetime import datetime @dataclass class Ticket: @@ -25,6 +26,33 @@ class Ticket: """Sets the ticket status to 'completed'.""" self.status = "completed" + def to_dict(self) -> Dict[str, Any]: + return { + "id": self.id, + "description": self.description, + "status": self.status, + "assigned_to": self.assigned_to, + "target_file": self.target_file, + "context_requirements": self.context_requirements, + "depends_on": self.depends_on, + "blocked_reason": self.blocked_reason, + "step_mode": self.step_mode, + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "Ticket": + return cls( + id=data["id"], + description=data.get("description"), + status=data.get("status"), + assigned_to=data.get("assigned_to"), + target_file=data.get("target_file"), + context_requirements=data.get("context_requirements", []), + depends_on=data.get("depends_on", []), + blocked_reason=data.get("blocked_reason"), + step_mode=data.get("step_mode", False), + ) + @dataclass class Track: """ @@ -67,3 +95,65 @@ class WorkerContext: ticket_id: str model_name: str messages: List[dict] + +@dataclass +class Metadata: + id: str + name: str + status: str + created_at: datetime + updated_at: datetime + + def to_dict(self) -> Dict[str, Any]: + return { + "id": self.id, + "name": self.name, + "status": self.status, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "Metadata": + return cls( + id=data["id"], + name=data["name"], + status=data.get("status"), + created_at=datetime.fromisoformat(data['created_at']) if data.get('created_at') else None, + updated_at=datetime.fromisoformat(data['updated_at']) if data.get('updated_at') else None, + ) + +@dataclass +class TrackState: + metadata: Metadata + discussion: List[Dict[str, Any]] + tasks: List[Ticket] + + def to_dict(self) -> Dict[str, Any]: + return { + "metadata": self.metadata.to_dict(), + "discussion": [ + { + k: v.isoformat() if isinstance(v, datetime) else v + for k, v in item.items() + } + for item in self.discussion + ], + "tasks": [task.to_dict() for task in self.tasks], + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "TrackState": + metadata = Metadata.from_dict(data["metadata"]) + tasks = [Ticket.from_dict(task_data) for task_data in data["tasks"]] + return cls( + metadata=metadata, + discussion=[ + { + k: datetime.fromisoformat(v) if isinstance(v, str) and 'T' in v else v # Basic check for ISO format + for k, v in item.items() + } + for item in data["discussion"] + ], + tasks=tasks, + ) diff --git a/tests/test_track_state_schema.py b/tests/test_track_state_schema.py new file mode 100644 index 0000000..9754ae0 --- /dev/null +++ b/tests/test_track_state_schema.py @@ -0,0 +1,172 @@ +import pytest +from datetime import datetime, timezone, timedelta + +# Import necessary classes from models.py +from models import Metadata, TrackState, Ticket + +# --- Pytest Tests --- + +def test_track_state_instantiation(): + """Test creating a TrackState object.""" + now = datetime.now(timezone.utc) + metadata = Metadata( + id="track-123", + name="Initial Setup", + status="in_progress", + created_at=now - timedelta(days=1), + updated_at=now, + ) + discussion = [ + {"role": "user", "content": "Hello", "ts": now - timedelta(hours=1)}, + {"role": "assistant", "content": "Hi there!", "ts": now - timedelta(hours=2)}, + ] + # Update Ticket instantiation to match models.py fields (description, assigned_to) + tasks = [ + Ticket(id="task-a", description="Design UI", status="todo", assigned_to="dev1"), + Ticket(id="task-b", description="Implement Backend", status="todo", assigned_to="dev2"), + ] + + track_state = TrackState( + metadata=metadata, + discussion=discussion, + tasks=tasks, + ) + + assert track_state.metadata.id == "track-123" + assert len(track_state.discussion) == 2 + assert len(track_state.tasks) == 2 + assert isinstance(track_state.tasks[0], Ticket) + assert track_state.tasks[0].description == "Design UI" + assert track_state.tasks[0].assigned_to == "dev1" + +def test_track_state_to_dict(): + """Test the to_dict() method for serialization.""" + now = datetime.now(timezone.utc) + metadata = Metadata( + id="track-456", + name="Refinement Phase", + status="completed", + created_at=now - timedelta(days=5), + updated_at=now - timedelta(days=2), + ) + discussion = [ + {"role": "user", "content": "Need changes", "ts": now - timedelta(hours=3)}, + {"role": "assistant", "content": "Understood.", "ts": now - timedelta(hours=4)}, + ] + # Update Ticket instantiation + tasks = [ + Ticket(id="task-c", description="Add feature X", status="in_progress", assigned_to="dev3"), + ] + + track_state = TrackState( + metadata=metadata, + discussion=discussion, + tasks=tasks, + ) + + track_dict = track_state.to_dict() + + assert track_dict["metadata"]["id"] == "track-456" + assert track_dict["metadata"]["created_at"] == metadata.created_at.isoformat() + assert track_dict["metadata"]["updated_at"] == metadata.updated_at.isoformat() + assert len(track_dict["discussion"]) == 2 + assert track_dict["discussion"][0]["ts"] == discussion[0]["ts"].isoformat() + assert len(track_dict["tasks"]) == 1 + # Use the Ticket's to_dict method for serialization + assert track_dict["tasks"][0]["id"] == "task-c" + assert track_dict["tasks"][0]["description"] == "Add feature X" + assert track_dict["tasks"][0]["assigned_to"] == "dev3" + +def test_track_state_from_dict(): + """Test the from_dict() class method for deserialization.""" + now = datetime.now(timezone.utc) + track_dict_data = { + "metadata": { + "id": "track-789", + "name": "Final Review", + "status": "pending", + "created_at": (now - timedelta(days=10)).isoformat(), + "updated_at": (now - timedelta(days=9)).isoformat(), + }, + "discussion": [ + {"role": "user", "content": "Review complete.", "ts": (now - timedelta(hours=5)).isoformat()}, + ], + "tasks": [ + # Use fields from models.py Ticket definition for deserialization + {"id": "task-d", "description": "Deploy", "status": "completed", "assigned_to": "ops1"}, + ], + } + + track_state = TrackState.from_dict(track_dict_data) + + assert isinstance(track_state, TrackState) + assert track_state.metadata.id == "track-789" + assert isinstance(track_state.metadata.created_at, datetime) + assert track_state.metadata.created_at.isoformat() == track_dict_data["metadata"]["created_at"] + assert len(track_state.discussion) == 1 + assert isinstance(track_state.discussion[0]["ts"], datetime) + assert track_state.discussion[0]["ts"].isoformat() == track_dict_data["discussion"][0]["ts"] + assert len(track_state.tasks) == 1 + assert isinstance(track_state.tasks[0], Ticket) + assert track_state.tasks[0].id == "task-d" + assert track_state.tasks[0].description == "Deploy" + assert track_state.tasks[0].assigned_to == "ops1" + +# Test case for empty lists and missing keys for robustness +def test_track_state_from_dict_empty_and_missing(): + """Test from_dict with empty lists and missing optional keys.""" + track_dict_data = { + "metadata": { + "id": "track-empty", + "name": "Empty State", + # created_at, updated_at, status are optional in from_dict logic + }, + "discussion": [], # Empty discussion list + "tasks": [], # Empty tasks list + } + + track_state = TrackState.from_dict(track_dict_data) + + assert isinstance(track_state, TrackState) + assert track_state.metadata.id == "track-empty" + assert track_state.metadata.name == "Empty State" + assert track_state.metadata.created_at is None + assert track_state.metadata.updated_at is None + assert track_state.metadata.status is None + assert len(track_state.discussion) == 0 + assert len(track_state.tasks) == 0 + +# Test case for to_dict with None values or missing optional data +def test_track_state_to_dict_with_none(): + """Test to_dict with None values in optional fields.""" + now = datetime.now(timezone.utc) + metadata = Metadata( + id="track-none", + name="None Test", + status=None, # None status + created_at=now, + updated_at=None, # None updated_at + ) + discussion = [ + {"role": "system", "content": "Info", "ts": None}, # None timestamp + ] + # Update Ticket instantiation + tasks = [ + Ticket(id="task-none", description="Task None", status="pending", assigned_to="anon"), + ] + + track_state = TrackState( + metadata=metadata, + discussion=discussion, + tasks=tasks, + ) + + track_dict = track_state.to_dict() + + assert track_dict["metadata"]["status"] is None + # Check that isoformat was called on datetime object, not None + assert track_dict["metadata"]["created_at"] == now.isoformat() + assert track_dict["metadata"]["updated_at"] is None # This should be None as it's passed as None + assert track_dict["discussion"][0]["ts"] is None + assert track_dict["tasks"][0]["description"] == "Task None" + assert track_dict["tasks"][0]["assigned_to"] == "anon"