core/__init__.py

@@ -2,20 +2,14 @@
 
 from core.models import (
     Bounty,
-    GroupBounty,
-    PersonalBounty,
     TrackedBounty,
-    GroupData,
+    RoomData,
     TrackingData,
-    UserData,
 )
 
 __all__ = [
     "Bounty",
-    "GroupBounty",
-    "PersonalBounty",
     "TrackedBounty",
-    "GroupData",
+    "RoomData",
     "TrackingData",
-    "UserData",
 ]

core/models.py

@@ -6,27 +6,18 @@ from typing import Optional
 
 @dataclass
 class Bounty:
-    """Base bounty with common fields."""
+    """Bounty - used for both group and personal bounties.
+
+    Use created_by_user_id to distinguish: if set, it's a group bounty
+    created by that user. If None, it's a personal/DM bounty.
+    """
 
     id: int
     text: Optional[str]
     link: Optional[str]
     due_date_ts: Optional[int]
     created_at: int
-
-
-@dataclass
-class GroupBounty(Bounty):
-    """Bounty created in a group context."""
-
-    created_by_user_id: int
-
-
-@dataclass
-class PersonalBounty(Bounty):
-    """Bounty created in DM/personal context."""
-
-    pass
+    created_by_user_id: Optional[int] = None
 
 
 @dataclass
@@ -38,26 +29,27 @@ class TrackedBounty:
 
 
 @dataclass
-class GroupData:
-    """All data for a group."""
+class RoomData:
+    """All data for a room (group or DM).
 
-    group_id: int
-    bounties: list[GroupBounty]
+    For groups: is_group=True, room_id is negative (Telegram group ID)
+    For DMs: is_group=False, room_id is the user_id (positive)
+    next_id is used to generate unique bounty IDs within this room.
+    """
+
+    room_id: int
+    is_group: bool
+    bounties: list[Bounty]
     next_id: int
 
 
 @dataclass
 class TrackingData:
-    """User tracking data within a group."""
+    """User tracking data within a group.
 
+    group_id is the Telegram group ID (always negative).
+    """
+
+    group_id: int
     user_id: int
     tracked: list[TrackedBounty]
-
-
-@dataclass
-class UserData:
-    """All personal bounties for a user (DM mode)."""
-
-    user_id: int
-    bounties: list[PersonalBounty]
-    next_id: int

tests/test_models.py

@@ -6,12 +6,9 @@ import pytest
 
 from core.models import (
     Bounty,
-    GroupBounty,
-    PersonalBounty,
     TrackedBounty,
-    GroupData,
+    RoomData,
     TrackingData,
-    UserData,
 )
 
 
@@ -23,12 +20,14 @@ class TestBounty:
             link="https://github.com/example/repo/issues/1",
             due_date_ts=1735689600,
             created_at=1735603200,
+            created_by_user_id=None,
         )
         assert b.id == 1
         assert b.text == "Fix the bug"
         assert b.link == "https://github.com/example/repo/issues/1"
         assert b.due_date_ts == 1735689600
         assert b.created_at == 1735603200
+        assert b.created_by_user_id is None
 
     def test_bounty_optional_fields_can_be_none(self):
         b = Bounty(
@@ -37,20 +36,52 @@ class TestBounty:
             link=None,
             due_date_ts=None,
             created_at=0,
+            created_by_user_id=None,
         )
         assert b.text is None
         assert b.link is None
         assert b.due_date_ts is None
 
-    def test_bounty_comparison(self):
-        b1 = Bounty(id=1, text="a", link=None, due_date_ts=None, created_at=0)
-        b2 = Bounty(id=1, text="a", link=None, due_date_ts=None, created_at=0)
+    def test_bounty_comparison_equal(self):
+        b1 = Bounty(
+            id=1,
+            text="a",
+            link=None,
+            due_date_ts=None,
+            created_at=0,
+            created_by_user_id=None,
+        )
+        b2 = Bounty(
+            id=1,
+            text="a",
+            link=None,
+            due_date_ts=None,
+            created_at=0,
+            created_by_user_id=None,
+        )
         assert b1 == b2
 
+    def test_bounty_comparison_not_equal(self):
+        b1 = Bounty(
+            id=1,
+            text="a",
+            link=None,
+            due_date_ts=None,
+            created_at=0,
+            created_by_user_id=None,
+        )
+        b2 = Bounty(
+            id=2,
+            text="b",
+            link=None,
+            due_date_ts=None,
+            created_at=0,
+            created_by_user_id=None,
+        )
+        assert b1 != b2
 
-class TestGroupBounty:
-    def test_create_group_bounty(self):
-        gb = GroupBounty(
+    def test_bounty_with_created_by_user_id(self):
+        b = Bounty(
             id=1,
             text="Group task",
             link=None,
@@ -58,43 +89,7 @@ class TestGroupBounty:
             created_at=0,
             created_by_user_id=123456,
         )
-        assert gb.id == 1
-        assert gb.text == "Group task"
-        assert gb.created_by_user_id == 123456
-
-    def test_group_bounty_inherits_from_bounty(self):
-        gb = GroupBounty(
-            id=1,
-            text="Task",
-            link="https://example.com",
-            due_date_ts=int(time.time()),
-            created_at=0,
-            created_by_user_id=999,
-        )
-        assert isinstance(gb, Bounty)
-
-
-class TestPersonalBounty:
-    def test_create_personal_bounty(self):
-        pb = PersonalBounty(
-            id=1,
-            text="Personal task",
-            link=None,
-            due_date_ts=None,
-            created_at=0,
-        )
-        assert pb.id == 1
-        assert pb.text == "Personal task"
-
-    def test_personal_bounty_inherits_from_bounty(self):
-        pb = PersonalBounty(
-            id=1,
-            text="Task",
-            link=None,
-            due_date_ts=None,
-            created_at=0,
-        )
-        assert isinstance(pb, Bounty)
+        assert b.created_by_user_id == 123456
 
 
 class TestTrackedBounty:
@@ -109,19 +104,33 @@ class TestTrackedBounty:
         assert tb1 == tb2
 
 
-class TestGroupData:
-    def test_create_group_data(self):
-        gd = GroupData(
-            group_id=-1001,
+class TestRoomData:
+    def test_create_group_room_data(self):
+        rd = RoomData(
+            room_id=-1001,
+            is_group=True,
             bounties=[],
             next_id=1,
         )
-        assert gd.group_id == -1001
-        assert gd.bounties == []
-        assert gd.next_id == 1
+        assert rd.room_id == -1001
+        assert rd.is_group is True
+        assert rd.bounties == []
+        assert rd.next_id == 1
 
-    def test_group_data_with_bounties(self):
-        gb = GroupBounty(
+    def test_create_dm_room_data(self):
+        rd = RoomData(
+            room_id=123456,
+            is_group=False,
+            bounties=[],
+            next_id=1,
+        )
+        assert rd.room_id == 123456
+        assert rd.is_group is False
+        assert rd.bounties == []
+        assert rd.next_id == 1
+
+    def test_room_data_with_bounties(self):
+        b = Bounty(
             id=1,
             text="Task",
             link=None,
@@ -129,39 +138,21 @@ class TestGroupData:
             created_at=0,
             created_by_user_id=123,
         )
-        gd = GroupData(group_id=-1001, bounties=[gb], next_id=2)
-        assert len(gd.bounties) == 1
-        assert gd.bounties[0].text == "Task"
+        rd = RoomData(room_id=-1001, is_group=True, bounties=[b], next_id=2)
+        assert len(rd.bounties) == 1
+        assert rd.bounties[0].text == "Task"
+        assert rd.bounties[0].created_by_user_id == 123
 
 
 class TestTrackingData:
     def test_create_tracking_data(self):
-        td = TrackingData(user_id=123456, tracked=[])
+        td = TrackingData(group_id=-1001, user_id=123456, tracked=[])
+        assert td.group_id == -1001
         assert td.user_id == 123456
         assert td.tracked == []
 
     def test_tracking_data_with_tracked(self):
         tb = TrackedBounty(bounty_id=5, created_at=0)
-        td = TrackingData(user_id=123, tracked=[tb])
+        td = TrackingData(group_id=-1001, user_id=123, tracked=[tb])
         assert len(td.tracked) == 1
         assert td.tracked[0].bounty_id == 5
-
-
-class TestUserData:
-    def test_create_user_data(self):
-        ud = UserData(user_id=123456, bounties=[], next_id=1)
-        assert ud.user_id == 123456
-        assert ud.bounties == []
-        assert ud.next_id == 1
-
-    def test_user_data_with_bounties(self):
-        pb = PersonalBounty(
-            id=1,
-            text="My task",
-            link=None,
-            due_date_ts=None,
-            created_at=0,
-        )
-        ud = UserData(user_id=123, bounties=[pb], next_id=2)
-        assert len(ud.bounties) == 1
-        assert ud.bounties[0].text == "My task"