Modernized codebase with NamedTuples, StrEnum, override decorators, slots, and other idiomatic improvements.

crabstero/bot.py

@@ -20,6 +20,7 @@ cog loading, ingestion worker pool, and graceful shutdown.
 
 import asyncio
 import logging
+from typing import override
 
 import discord
 from discord import app_commands
@@ -41,6 +42,8 @@ from crabstero.tasks.ingestion import ingest_channel
 
 logger = logging.getLogger(__name__)
 
+type IngestableChannel = discord.TextChannel | discord.VoiceChannel
+
 
 class Crabstero(commands.Bot):
     """Central bot subclass that owns all lifecycle state.
@@ -80,9 +83,7 @@ class Crabstero(commands.Bot):
         self._database_path = database_path
         self._ingestion_worker_count = ingestion_workers
         self._ingest_only = ingest_only
-        self._ingestion_queue: asyncio.Queue[
+        self._ingestion_queue: asyncio.Queue[IngestableChannel] = asyncio.Queue()
-            discord.TextChannel | discord.VoiceChannel
-        ] = asyncio.Queue()
         self._ingestion_workers: list[asyncio.Task[None]] = []
         self._db: Database | None = None
         self.ingest_cache = IngestCache()
@@ -92,7 +93,7 @@ class Crabstero(commands.Bot):
 
         DISCORD_LATENCY.set_function(lambda: self.latency)
         GUILD_COUNT.set_function(lambda: len(self.guilds))
-        INGESTION_BACKLOG.set_function(lambda: self._ingestion_queue.qsize())
+        INGESTION_BACKLOG.set_function(self._ingestion_queue.qsize)
 
         repo_url = "https://git.logal.dev/LogalDeveloper/Crabstero"
         self.http.user_agent = f"DiscordBot ({repo_url}, {crabstero_version})"
@@ -112,6 +113,7 @@ class Crabstero(commands.Bot):
         """Whether the bot is running in ingest-only mode."""
         return self._ingest_only
 
+    @override
     async def setup_hook(self) -> None:
         """Open the database, start ingestion workers, and load all cogs."""
         self._db = await Database.connect(self._database_path)
@@ -149,6 +151,7 @@ class Crabstero(commands.Bot):
                 logger.info("Slash command tree has changed, syncing with Discord.")
                 await self.tree.sync()
 
+    @override
     def dispatch(self, event: str, /, *args: object, **kwargs: object) -> None:
         """Dispatch an event, incrementing the events counter.
 
@@ -161,10 +164,12 @@ class Crabstero(commands.Bot):
         """Log that the bot has started successfully."""
         logger.info("Crabstero started!")
 
+    @override
     async def start(self, token: str = "", *, reconnect: bool = True) -> None:
         """Start the bot using the token provided at initialization."""
         await super().start(self._token, reconnect=reconnect)
 
+    @override
     async def close(self) -> None:
         """Cancel ingestion workers, close the database, and then the bot connection."""
         if self.is_closed():
@@ -181,9 +186,7 @@ class Crabstero(commands.Bot):
             await self._db.close()
         await super().close()
 
-    def queue_channel_for_ingestion(
+    def queue_channel_for_ingestion(self, channel: IngestableChannel) -> None:
-        self, channel: discord.TextChannel | discord.VoiceChannel
-    ) -> None:
         """Enqueue a single channel for background message history ingestion.
 
         :param channel: The channel to enqueue.

crabstero/cache.py

@@ -47,6 +47,8 @@ class IngestCache:
     :param cleanup_interval_seconds: How often the background task runs.
     """
 
+    __slots__ = ("_cleanup_interval", "_entries", "_task", "_ttl")
+
     def __init__(
         self,
         ttl_seconds: float = 120,

crabstero/database.py

@@ -14,10 +14,36 @@
 
 """Async SQLite database access for Crabstero's persistent storage."""
 
-from typing import Self
+from typing import NamedTuple, Self
 
 import aiosqlite
 
+
+class StartWord(NamedTuple):
+    """A Markov chain starting word row."""
+
+    channel_id: int
+    user_id: int
+    word: str
+
+
+class Transition(NamedTuple):
+    """A Markov chain word transition row."""
+
+    channel_id: int
+    user_id: int
+    word: str
+    next_word: str
+
+
+class ChannelImage(NamedTuple):
+    """An image URL associated with a channel."""
+
+    channel_id: int
+    user_id: int
+    url: str
+
+
 # SQL statements for creating the database schema.
 _SCHEMA = """
 -- Markov chain starting words.
@@ -107,15 +133,15 @@ class Database:
 
     async def add_markov_data(
         self,
-        start_words: list[tuple[int, int, str]],
+        start_words: list[StartWord],
-        transitions: list[tuple[int, int, str, str]],
+        transitions: list[Transition],
     ) -> None:
         """Insert Markov start words and transitions, then commit.
 
         Both inserts happen in a single transaction.
 
-        :param start_words: A list of (channel_id, user_id, word) tuples.
+        :param start_words: Starting word entries to insert.
-        :param transitions: A list of (channel_id, user_id, word, next_word) tuples.
+        :param transitions: Transition entries to insert.
         """
         if start_words:
             await self._connection.executemany(
@@ -136,16 +162,16 @@ class Database:
 
     async def remove_markov_data(
         self,
-        start_words: list[tuple[int, int, str]],
+        start_words: list[StartWord],
-        transitions: list[tuple[int, int, str, str]],
+        transitions: list[Transition],
     ) -> None:
         """Remove one matching row per entry from the Markov tables.
 
         Each entry removes at most one duplicate row, preserving remaining
         frequency weight.
 
-        :param start_words: A list of (channel_id, user_id, word) tuples.
+        :param start_words: Starting word entries to remove.
-        :param transitions: A list of (channel_id, user_id, word, next_word) tuples.
+        :param transitions: Transition entries to remove.
         """
         for channel_id, user_id, word in start_words:
             await self._connection.execute(
@@ -213,6 +239,8 @@ class Database:
         Weighted by occurrence frequency. A completing word is one whose last
         character is '.', '!', '?', or '§'.
 
+        The set of terminators must match ``markov._TERMINATORS``.
+
         :param channel_id: The Discord channel ID.
         :param word: The current word to find a completing transition for.
         :return: A random completing next word, or None if none exist.
@@ -228,10 +256,10 @@ class Database:
             row = await cursor.fetchone()
             return row[0] if row else None
 
-    async def add_images(self, images: list[tuple[int, int, str]]) -> None:
+    async def add_images(self, images: list[ChannelImage]) -> None:
         """Store image URLs for a given channel.
 
-        :param images: A list of (channel_id, user_id, url) tuples.
+        :param images: Image entries to insert.
         """
         if images:
             await self._connection.executemany(
@@ -242,13 +270,13 @@ class Database:
             )
             await self._connection.commit()
 
-    async def remove_images(self, images: list[tuple[int, int, str]]) -> None:
+    async def remove_images(self, images: list[ChannelImage]) -> None:
         """Remove one matching row per entry from the images table.
 
         Each entry removes at most one duplicate row, preserving remaining
         frequency weight.
 
-        :param images: A list of (channel_id, user_id, url) tuples.
+        :param images: Image entries to remove.
         """
         for channel_id, user_id, url in images:
             await self._connection.execute(

crabstero/flags.py

@@ -27,7 +27,7 @@ if TYPE_CHECKING:
     from crabstero.database import Database
 
 
-class Flag(enum.Enum):
+class Flag(enum.StrEnum):
     """Enum of all supported flag names."""
 
     NO_REPLY = "noReply"
@@ -35,7 +35,7 @@ class Flag(enum.Enum):
     ALLOW_PINGS = "allowPings"
 
 
-class EntityType(enum.Enum):
+class EntityType(enum.StrEnum):
     """Enum of entity types that can have flags."""
 
     CHANNEL = "channel"
@@ -65,7 +65,7 @@ async def set_flag(
     :param entity_type: The type of the entity.
     :param flag: The flag to set.
     """
-    await db.set_flag(entity_type.value, _entity_id(entity), flag.value)
+    await db.set_flag(entity_type, _entity_id(entity), flag)
 
 
 async def clear_flag(
@@ -81,7 +81,7 @@ async def clear_flag(
     :param entity_type: The type of the entity.
     :param flag: The flag to clear.
     """
-    await db.clear_flag(entity_type.value, _entity_id(entity), flag.value)
+    await db.clear_flag(entity_type, _entity_id(entity), flag)
 
 
 async def is_flag_set(
@@ -98,4 +98,4 @@ async def is_flag_set(
     :param flag: The flag to check for.
     :return: True if the flag is set, False otherwise.
     """
-    return await db.is_flag_set(entity_type.value, _entity_id(entity), flag.value)
+    return await db.is_flag_set(entity_type, _entity_id(entity), flag)

crabstero/markov.py

@@ -24,6 +24,7 @@ import re
 from typing import TYPE_CHECKING
 
 from crabstero import metrics
+from crabstero.database import StartWord, Transition
 
 if TYPE_CHECKING:
     from crabstero.database import Database
@@ -116,8 +117,8 @@ async def _ingest_sentence(
     :param sentence: The sentence to ingest.
     """
     raw_starts, raw_transitions = _tokenize_sentence(sentence)
-    start_words = [(channel_id, user_id, w) for w in raw_starts]
+    start_words = [StartWord(channel_id, user_id, w) for w in raw_starts]
-    transitions = [(channel_id, user_id, w, nw) for w, nw in raw_transitions]
+    transitions = [Transition(channel_id, user_id, w, nw) for w, nw in raw_transitions]
     await db.add_markov_data(start_words, transitions)
 
 
@@ -147,8 +148,8 @@ async def _uningest_sentence(
     :param sentence: The sentence to uningest.
     """
     raw_starts, raw_transitions = _tokenize_sentence(sentence)
-    start_words = [(channel_id, user_id, w) for w in raw_starts]
+    start_words = [StartWord(channel_id, user_id, w) for w in raw_starts]
-    transitions = [(channel_id, user_id, w, nw) for w, nw in raw_transitions]
+    transitions = [Transition(channel_id, user_id, w, nw) for w, nw in raw_transitions]
     await db.remove_markov_data(start_words, transitions)
 
 
@@ -201,8 +202,7 @@ async def generate(
             result = result[:hard_limit]
 
         # Strip the internal sentence-end marker so it never appears in output.
-        if result.endswith(DEFAULT_SENTENCE_END):
+        result = result.removesuffix(DEFAULT_SENTENCE_END)
-            result = result[:-1]
 
         metrics.GENERATED_MESSAGE_LENGTH.observe(len(result))
         return result

crabstero/messages.py

@@ -27,14 +27,15 @@ import discord
 
 from crabstero import flags, markov, metrics
 from crabstero.cache import CachedMessage, IngestCache
+from crabstero.database import ChannelImage
 from crabstero.flags import EntityType, Flag
 
 if TYPE_CHECKING:
     from crabstero.database import Database
 
 # Copied from discordjs/discord-api-types:
-# https://github.com/discordjs/discord-api-types/blob/7fe434114e91c80ed79f0204ae6c73047672d55d/globals.ts#L30
+# https://github.com/discordjs/discord-api-types/blob/662cb0cb0ac9c6f9ad93e180849476714bfceb0c/globals.ts#L39
-MENTION_PATTERN = re.compile(r"<@!?(?P<id>\d{17,20})>")
+_MENTION_PATTERN = re.compile(r"<@!?(?P<id>\d{17,20})>")
 
 _EMBED_CHANCE_THRESHOLD = 95  # Out of 100; sends an embed ~5% of the time.
 
@@ -89,7 +90,7 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
     # Suppress all mentions by default; only ping users who opted in.
     allowed_user_ids: set[int] = set()
 
-    for match in MENTION_PATTERN.finditer(body):
+    for match in _MENTION_PATTERN.finditer(body):
         user_id = int(match.group("id"))
 
         if await flags.is_flag_set(db, user_id, EntityType.USER, Flag.ALLOW_PINGS):
@@ -169,7 +170,9 @@ async def ingest_message(
                 image_urls.append(embed.image.url)
 
         if image_urls:
-            await db.add_images([(channel_id, user_id, url) for url in image_urls])
+            await db.add_images(
+                [ChannelImage(channel_id, user_id, url) for url in image_urls]
+            )
 
     metrics.MESSAGES_INGESTED.inc()
 
@@ -208,7 +211,10 @@ async def uningest_message(db: Database, cache: IngestCache, message_id: int) ->
 
     if entry.image_urls:
         await db.remove_images(
-            [(entry.channel_id, entry.user_id, url) for url in entry.image_urls]
+            [
+                ChannelImage(entry.channel_id, entry.user_id, url)
+                for url in entry.image_urls
+            ]
         )
 
     metrics.MESSAGES_UNINGESTED.inc()

crabstero/metrics.py

@@ -131,6 +131,8 @@ class MetricsServer:
     which handles compression and content negotiation automatically.
     """
 
+    __slots__ = ("_host", "_port", "_runner")
+
     def __init__(self, host: str, port: int) -> None:
         """Store the listen address.
 

crabstero/tasks/ingestion.py

@@ -26,7 +26,7 @@ from crabstero import metrics
 from crabstero.messages import ingest_message
 
 if TYPE_CHECKING:
-    from crabstero.bot import Crabstero
+    from crabstero.bot import Crabstero, IngestableChannel
     from crabstero.database import Database
 
 MAXIMUM_MESSAGES_PER_CHANNEL = (
@@ -48,7 +48,7 @@ def queue_channels_for_ingestion(guild: discord.Guild, bot: Crabstero) -> None:
 
 
 async def ingest_channel(
-    channel: discord.TextChannel | discord.VoiceChannel,
+    channel: IngestableChannel,
     db: Database,
 ) -> None:
     """Bulk-ingest the message history of a given channel.

pyproject.toml

@@ -113,8 +113,6 @@ ignore = [
 asyncio_mode = "auto"
 
 [tool.ruff.lint.per-file-ignores]
-"crabstero/__init__.py" = ["E402"]     # constants defined before imports intentionally
-"crabstero/__main__.py" = ["T201"]     # CLI entry point uses print() for user output
 "tests/**" = ["S101"]                  # assert is standard for pytest
 
 [tool.coverage.run]

tests/test_database.py

@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING
 
 import pytest
 
-from crabstero.database import Database
+from crabstero.database import ChannelImage, Database, StartWord, Transition
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -58,21 +58,21 @@ class TestAddMarkovData:
 
     async def test_stores_start_word(self, db: Database) -> None:
         """Inserted start word can be retrieved by channel."""
-        await db.add_markov_data([(1, 100, "Hello")], [])
+        await db.add_markov_data([StartWord(1, 100, "Hello")], [])
         result = await db.get_random_start_word(1)
         assert result == "Hello"
 
     async def test_stores_transition(self, db: Database) -> None:
         """Inserted transition can be retrieved by channel and word."""
-        await db.add_markov_data([], [(1, 100, "Hello", "world.")])
+        await db.add_markov_data([], [Transition(1, 100, "Hello", "world.")])
         result = await db.get_random_next_word(1, "Hello")
         assert result == "world."
 
     async def test_stores_both_in_single_call(self, db: Database) -> None:
         """Start words and transitions are stored in a single call."""
         await db.add_markov_data(
-            [(1, 100, "Hello")],
+            [StartWord(1, 100, "Hello")],
-            [(1, 100, "Hello", "world.")],
+            [Transition(1, 100, "Hello", "world.")],
         )
         assert await db.get_random_start_word(1) == "Hello"
         assert await db.get_random_next_word(1, "Hello") == "world."
@@ -98,8 +98,8 @@ class TestAddMarkovData:
         await db.add_markov_data(
             [],
             [
-                (1, 100, "Hello", "beautiful"),
+                Transition(1, 100, "Hello", "beautiful"),
-                (1, 100, "Hello", completing_word),
+                Transition(1, 100, "Hello", completing_word),
             ],
         )
         for _ in range(100):
@@ -108,7 +108,7 @@ class TestAddMarkovData:
 
     async def test_completing_returns_none_without_match(self, db: Database) -> None:
         """Returns None when no transitions end with sentence punctuation."""
-        await db.add_markov_data([], [(1, 100, "Hello", "beautiful")])
+        await db.add_markov_data([], [Transition(1, 100, "Hello", "beautiful")])
         result = await db.get_random_completing_next_word(1, "Hello")
         assert result is None
 
@@ -127,7 +127,7 @@ class TestImages:
 
     async def test_add_and_retrieve(self, db: Database) -> None:
         """Inserted image URL can be retrieved by channel."""
-        await db.add_images([(1, 100, "https://example.com/cat.png")])
+        await db.add_images([ChannelImage(1, 100, "https://example.com/cat.png")])
         result = await db.get_random_image(1)
         assert result == "https://example.com/cat.png"
 
@@ -186,48 +186,63 @@ class TestRemoveMarkovData:
 
     async def test_removes_one_start_word(self, db: Database) -> None:
         """Removes exactly one matching start word row."""
-        await db.add_markov_data([(1, 100, "Hello"), (1, 100, "Hello")], [])
+        await db.add_markov_data(
-        await db.remove_markov_data([(1, 100, "Hello")], [])
+            [StartWord(1, 100, "Hello"), StartWord(1, 100, "Hello")], []
+        )
+        await db.remove_markov_data([StartWord(1, 100, "Hello")], [])
         # One copy should remain.
         assert await db.get_random_start_word(1) == "Hello"
 
     async def test_removes_one_transition(self, db: Database) -> None:
         """Removes exactly one matching transition row."""
         await db.add_markov_data(
-            [], [(1, 100, "Hello", "world."), (1, 100, "Hello", "world.")]
+            [],
+            [
+                Transition(1, 100, "Hello", "world."),
+                Transition(1, 100, "Hello", "world."),
+            ],
         )
-        await db.remove_markov_data([], [(1, 100, "Hello", "world.")])
+        await db.remove_markov_data([], [Transition(1, 100, "Hello", "world.")])
         assert await db.get_random_next_word(1, "Hello") == "world."
 
     async def test_removes_last_start_word(self, db: Database) -> None:
         """Removing the only start word leaves the table empty for that channel."""
-        await db.add_markov_data([(1, 100, "Hello")], [])
+        await db.add_markov_data([StartWord(1, 100, "Hello")], [])
-        await db.remove_markov_data([(1, 100, "Hello")], [])
+        await db.remove_markov_data([StartWord(1, 100, "Hello")], [])
         assert await db.get_random_start_word(1) is None
 
     async def test_removes_last_transition(self, db: Database) -> None:
         """Removing the only transition leaves no next word."""
-        await db.add_markov_data([], [(1, 100, "Hello", "world.")])
+        await db.add_markov_data([], [Transition(1, 100, "Hello", "world.")])
-        await db.remove_markov_data([], [(1, 100, "Hello", "world.")])
+        await db.remove_markov_data([], [Transition(1, 100, "Hello", "world.")])
         assert await db.get_random_next_word(1, "Hello") is None
 
     async def test_no_match_is_noop(self, db: Database) -> None:
         """Removing a non-existent row does not raise."""
-        await db.remove_markov_data([(1, 100, "nope")], [(1, 100, "nope", "nah")])
+        await db.remove_markov_data(
+            [StartWord(1, 100, "nope")],
+            [Transition(1, 100, "nope", "nah")],
+        )
 
     async def test_start_word_removal_scoped_to_channel(self, db: Database) -> None:
         """Removing a start word in one channel leaves another channel intact."""
-        await db.add_markov_data([(1, 100, "Hello"), (2, 200, "Hello")], [])
+        await db.add_markov_data(
-        await db.remove_markov_data([(1, 100, "Hello")], [])
+            [StartWord(1, 100, "Hello"), StartWord(2, 200, "Hello")], []
+        )
+        await db.remove_markov_data([StartWord(1, 100, "Hello")], [])
         assert await db.get_random_start_word(1) is None
         assert await db.get_random_start_word(2) == "Hello"
 
     async def test_transition_removal_scoped_to_channel(self, db: Database) -> None:
         """Removing a transition in one channel leaves another channel intact."""
         await db.add_markov_data(
-            [], [(1, 100, "Hello", "world."), (2, 200, "Hello", "world.")]
+            [],
+            [
+                Transition(1, 100, "Hello", "world."),
+                Transition(2, 200, "Hello", "world."),
+            ],
         )
-        await db.remove_markov_data([], [(1, 100, "Hello", "world.")])
+        await db.remove_markov_data([], [Transition(1, 100, "Hello", "world.")])
         assert await db.get_random_next_word(1, "Hello") is None
         assert await db.get_random_next_word(2, "Hello") == "world."
 
@@ -243,35 +258,35 @@ class TestRemoveImages:
         """Removes exactly one matching image row."""
         await db.add_images(
             [
-                (1, 100, "https://example.com/a.png"),
+                ChannelImage(1, 100, "https://example.com/a.png"),
-                (1, 100, "https://example.com/a.png"),
+                ChannelImage(1, 100, "https://example.com/a.png"),
             ]
         )
-        await db.remove_images([(1, 100, "https://example.com/a.png")])
+        await db.remove_images([ChannelImage(1, 100, "https://example.com/a.png")])
         # One copy should remain.
         assert await db.get_random_image(1) == "https://example.com/a.png"
 
     async def test_removes_last_image(self, db: Database) -> None:
         """Removing the only image leaves none for that channel."""
-        await db.add_images([(1, 100, "https://example.com/a.png")])
+        await db.add_images([ChannelImage(1, 100, "https://example.com/a.png")])
-        await db.remove_images([(1, 100, "https://example.com/a.png")])
+        await db.remove_images([ChannelImage(1, 100, "https://example.com/a.png")])
         assert await db.get_random_image(1) is None
 
     async def test_image_removal_scoped_to_channel(self, db: Database) -> None:
         """Removing an image in one channel leaves another channel intact."""
         await db.add_images(
             [
-                (1, 100, "https://example.com/a.png"),
+                ChannelImage(1, 100, "https://example.com/a.png"),
-                (2, 200, "https://example.com/a.png"),
+                ChannelImage(2, 200, "https://example.com/a.png"),
             ]
         )
-        await db.remove_images([(1, 100, "https://example.com/a.png")])
+        await db.remove_images([ChannelImage(1, 100, "https://example.com/a.png")])
         assert await db.get_random_image(1) is None
         assert await db.get_random_image(2) == "https://example.com/a.png"
 
     async def test_no_match_is_noop(self, db: Database) -> None:
         """Removing a non-existent image does not raise."""
-        await db.remove_images([(1, 100, "https://example.com/nope.png")])
+        await db.remove_images([ChannelImage(1, 100, "https://example.com/nope.png")])
 
     async def test_empty_list_is_noop(self, db: Database) -> None:
         """Empty list does not error."""
@@ -285,7 +300,10 @@ class TestWriteDurability:
         """Data written via add_markov_data is durable after close/reopen."""
         db_path = str(tmp_path / "durability.db")
         db = await Database.connect(db_path)
-        await db.add_markov_data([(1, 100, "Hello")], [(1, 100, "Hello", "world.")])
+        await db.add_markov_data(
+            [StartWord(1, 100, "Hello")],
+            [Transition(1, 100, "Hello", "world.")],
+        )
         await db.close()
 
         db2 = await Database.connect(db_path)

tests/test_markov.py

@@ -22,6 +22,7 @@ from typing import TYPE_CHECKING
 
 import pytest
 
+from crabstero.database import StartWord, Transition
 from crabstero.markov import (
     DEFAULT_SENTENCE_END,
     _ingest_sentence,
@@ -211,16 +212,16 @@ class TestGenerate:
         # continuing ("E") and completing ("end.") transition, so
         # get_random_completing_next_word deterministically picks "end.".
         await db.add_markov_data(
-            [(1, 100, "A")],
+            [StartWord(1, 100, "A")],
             [
-                (1, 100, "A", "B"),
+                Transition(1, 100, "A", "B"),
-                (1, 100, "B", "C"),
+                Transition(1, 100, "B", "C"),
-                (1, 100, "C", "D"),
+                Transition(1, 100, "C", "D"),
-                (1, 100, "D", "E"),
+                Transition(1, 100, "D", "E"),
-                (1, 100, "D", "F"),
+                Transition(1, 100, "D", "F"),
-                (1, 100, "D", "G"),
+                Transition(1, 100, "D", "G"),
-                (1, 100, "D", "H"),
+                Transition(1, 100, "D", "H"),
-                (1, 100, "D", "end."),
+                Transition(1, 100, "D", "end."),
             ],
         )
 
@@ -230,14 +231,14 @@ class TestGenerate:
 
     async def test_start_word_already_ends_sentence(self, db: Database) -> None:
         """Generation stops immediately when the start word is sentence-ending."""
-        await db.add_markov_data([(1, 100, "Yes.")], [])
+        await db.add_markov_data([StartWord(1, 100, "Yes.")], [])
 
         result = await generate(db, 1)
         assert result == "Yes."
 
     async def test_chain_dead_end(self, db: Database) -> None:
         """Generation stops when no next word exists (dead-end chain)."""
-        await db.add_markov_data([(1, 100, "Hello")], [])
+        await db.add_markov_data([StartWord(1, 100, "Hello")], [])
 
         # "Hello" has no transitions, so the loop breaks immediately.
         result = await generate(db, 1)
@@ -247,10 +248,10 @@ class TestGenerate:
         """Section sign at the truncation boundary is stripped."""
         # Build a chain: "A" -> "B" -> "C§".
         await db.add_markov_data(
-            [(1, 100, "A")],
+            [StartWord(1, 100, "A")],
             [
-                (1, 100, "A", "B"),
+                Transition(1, 100, "A", "B"),
-                (1, 100, "B", f"C{DEFAULT_SENTENCE_END}"),
+                Transition(1, 100, "B", f"C{DEFAULT_SENTENCE_END}"),
             ],
         )
 
@@ -264,14 +265,19 @@ class TestGenerate:
         """After soft_limit, falls back when no completing word exists."""
         # "A" -> "B" (no completing transition). Past soft_limit,
         # get_random_completing_next_word returns None, falls back to "B".
-        await db.add_markov_data([(1, 100, "A")], [(1, 100, "A", "B")])
+        await db.add_markov_data(
+            [StartWord(1, 100, "A")], [Transition(1, 100, "A", "B")]
+        )
 
         result = await generate(db, 1, soft_limit=1, hard_limit=1000)
         assert result == "A B"
 
     async def test_hard_limit_truncates_mid_word(self, db: Database) -> None:
         """Hard limit slices output even when it falls inside a word."""
-        await db.add_markov_data([(1, 100, "AB")], [(1, 100, "AB", "CDEF")])
+        await db.add_markov_data(
+            [StartWord(1, 100, "AB")],
+            [Transition(1, 100, "AB", "CDEF")],
+        )
 
         # "AB CDEF" is 7 chars; hard_limit=5 truncates to "AB CD".
         result = await generate(db, 1, soft_limit=100, hard_limit=5)

tests/test_uningest.py

@@ -19,6 +19,7 @@ from typing import TYPE_CHECKING
 import pytest
 
 from crabstero.cache import CachedMessage, IngestCache
+from crabstero.database import ChannelImage
 from crabstero.markov import ingest, uningest
 
 if TYPE_CHECKING:
@@ -50,8 +51,8 @@ class TestIngestUningestCycle:
 
     async def test_image_round_trip(self, db: Database) -> None:
         """Ingest and uningest an image leaves the database clean."""
-        await db.add_images([(1, 100, "https://example.com/cat.png")])
+        await db.add_images([ChannelImage(1, 100, "https://example.com/cat.png")])
-        await db.remove_images([(1, 100, "https://example.com/cat.png")])
+        await db.remove_images([ChannelImage(1, 100, "https://example.com/cat.png")])
 
         assert await db.get_random_image(1) is None
 
@@ -140,7 +141,7 @@ class TestUningestRestoresState:
     ) -> None:
         """Ingest then uningest preserves unrelated pre-existing data exactly."""
         await ingest(db, 99, 200, "Pre-existing data stays safe.")
-        await db.add_images([(99, 200, "https://example.com/existing.png")])
+        await db.add_images([ChannelImage(99, 200, "https://example.com/existing.png")])
 
         before = await _snapshot(db)
         await ingest(db, 1, 100, text)