Simplified database layer and consolidated Markov writes into a single transaction.
CI / Formatting (push) Successful in 5s
CI / Linting (push) Successful in 5s
CI / Tests (push) Successful in 16s
CI / Type Checking (push) Successful in 10s
CI / Spelling (push) Successful in 5s

This commit is contained in:
2026-03-17 17:07:28 -04:00
parent 56d3c36c7d
commit 6aa8527adc
5 changed files with 109 additions and 219 deletions
+33 -87
View File
@@ -12,25 +12,12 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Async SQLite database access for Crabstero's persistent storage. """Async SQLite database access for Crabstero's persistent storage."""
Uses aiosqlite for native async access. All methods are async def.
"""
import asyncio
import contextlib
import logging
from typing import Self from typing import Self
import aiosqlite import aiosqlite
logger = logging.getLogger(__name__)
AUTO_COMMIT_WRITE_THRESHOLD = 100 # Commit after this many DB write operations.
AUTO_COMMIT_TIMEOUT_SECONDS = (
60.0 # Commit after this many seconds since first uncommitted write.
)
# SQL statements for creating the database schema. # SQL statements for creating the database schema.
_SCHEMA = """ _SCHEMA = """
-- Markov chain starting words. -- Markov chain starting words.
@@ -83,8 +70,7 @@ class Database:
"""Manages all SQLite database operations for Crabstero. """Manages all SQLite database operations for Crabstero.
Uses aiosqlite for native async access. A single connection is held Uses aiosqlite for native async access. A single connection is held
open for the lifetime of the bot process with WAL mode enabled for open for the lifetime of the bot process.
concurrent read performance.
""" """
def __init__(self, connection: aiosqlite.Connection) -> None: def __init__(self, connection: aiosqlite.Connection) -> None:
@@ -93,8 +79,6 @@ class Database:
:param connection: An open aiosqlite connection. :param connection: An open aiosqlite connection.
""" """
self._connection = connection self._connection = connection
self._pending_writes = 0
self._flush_task: asyncio.Task[None] | None = None
@classmethod @classmethod
async def connect(cls, path: str) -> Self: async def connect(cls, path: str) -> Self:
@@ -108,9 +92,6 @@ class Database:
""" """
connection = await aiosqlite.connect(path) connection = await aiosqlite.connect(path)
# Enable WAL mode for better concurrent read performance.
await connection.execute("PRAGMA journal_mode=WAL")
# Set synchronous to NORMAL for a balance between safety and speed. # Set synchronous to NORMAL for a balance between safety and speed.
await connection.execute("PRAGMA synchronous=NORMAL") await connection.execute("PRAGMA synchronous=NORMAL")
@@ -120,38 +101,38 @@ class Database:
return cls(connection) return cls(connection)
async def commit(self) -> None:
"""Commit pending writes and reset the flush timer."""
await self._connection.commit()
self._pending_writes = 0
if self._flush_task is not None:
self._flush_task.cancel()
self._flush_task = None
async def close(self) -> None: async def close(self) -> None:
"""Cancel the flush timer, commit pending writes, and close the connection.""" """Close the database connection."""
if self._flush_task is not None:
self._flush_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self._flush_task
self._flush_task = None
if self._pending_writes > 0:
await self._connection.commit()
self._pending_writes = 0
await self._connection.close() await self._connection.close()
async def add_start_words_batch(self, rows: list[tuple[int, int, str]]) -> None: async def add_markov_data(
"""Insert a batch of starting words into the markov_start_words table. self,
start_words: list[tuple[int, int, str]],
transitions: list[tuple[int, int, str, str]],
) -> None:
"""Insert Markov start words and transitions, then commit.
:param rows: A list of (channel_id, user_id, word) tuples to insert. Both inserts happen in a single transaction.
:param start_words: A list of (channel_id, user_id, word) tuples.
:param transitions: A list of (channel_id, user_id, word, next_word) tuples.
""" """
await self._connection.executemany( if start_words:
"INSERT INTO markov_start_words" await self._connection.executemany(
" (channel_id, user_id, word)" "INSERT INTO markov_start_words"
" VALUES (?, ?, ?)", " (channel_id, user_id, word)"
rows, " VALUES (?, ?, ?)",
) start_words,
await self._maybe_commit() )
if transitions:
await self._connection.executemany(
"INSERT INTO markov_transitions"
" (channel_id, user_id, word, next_word)"
" VALUES (?, ?, ?, ?)",
transitions,
)
if start_words or transitions:
await self._connection.commit()
async def get_random_start_word(self, channel_id: int) -> str | None: async def get_random_start_word(self, channel_id: int) -> str | None:
"""Return a random starting word for a channel. """Return a random starting word for a channel.
@@ -170,22 +151,6 @@ class Database:
row = await cursor.fetchone() row = await cursor.fetchone()
return row[0] if row else None return row[0] if row else None
async def add_transitions_batch(
self, rows: list[tuple[int, int, str, str]]
) -> None:
"""Insert a batch of word transitions into the markov_transitions table.
:param rows: A list of (channel_id, user_id, word, next_word)
tuples to insert.
"""
await self._connection.executemany(
"INSERT INTO markov_transitions"
" (channel_id, user_id, word, next_word)"
" VALUES (?, ?, ?, ?)",
rows,
)
await self._maybe_commit()
async def get_random_next_word(self, channel_id: int, word: str) -> str | None: async def get_random_next_word(self, channel_id: int, word: str) -> str | None:
"""Return a random next word for a given word in a channel. """Return a random next word for a given word in a channel.
@@ -238,7 +203,7 @@ class Database:
"INSERT INTO channel_images (channel_id, user_id, url) VALUES (?, ?, ?)", "INSERT INTO channel_images (channel_id, user_id, url) VALUES (?, ?, ?)",
(channel_id, user_id, url), (channel_id, user_id, url),
) )
await self._maybe_commit() await self._connection.commit()
async def get_random_image(self, channel_id: int) -> str | None: async def get_random_image(self, channel_id: int) -> str | None:
"""Return a random image URL for a given channel. """Return a random image URL for a given channel.
@@ -268,7 +233,7 @@ class Database:
" VALUES (?, ?, ?)", " VALUES (?, ?, ?)",
(entity_type, entity_id, flag_name), (entity_type, entity_id, flag_name),
) )
await self._maybe_commit() await self._connection.commit()
async def clear_flag( async def clear_flag(
self, entity_type: str, entity_id: str, flag_name: str self, entity_type: str, entity_id: str, flag_name: str
@@ -286,7 +251,7 @@ class Database:
" AND flag_name = ?", " AND flag_name = ?",
(entity_type, entity_id, flag_name), (entity_type, entity_id, flag_name),
) )
await self._maybe_commit() await self._connection.commit()
async def is_flag_set( async def is_flag_set(
self, entity_type: str, entity_id: str, flag_name: str self, entity_type: str, entity_id: str, flag_name: str
@@ -328,23 +293,4 @@ class Database:
"INSERT OR IGNORE INTO ingested_channels (channel_id) VALUES (?)", "INSERT OR IGNORE INTO ingested_channels (channel_id) VALUES (?)",
(channel_id,), (channel_id,),
) )
await self._maybe_commit() await self._connection.commit()
async def _maybe_commit(self) -> None:
"""Track a pending write and commit or start a flush timer."""
self._pending_writes += 1
if self._pending_writes >= AUTO_COMMIT_WRITE_THRESHOLD:
await self.commit()
elif self._flush_task is None:
self._flush_task = asyncio.create_task(self._flush_after_timeout())
async def _flush_after_timeout(self) -> None:
"""Background task that commits after the timeout elapses."""
try:
await asyncio.sleep(AUTO_COMMIT_TIMEOUT_SECONDS)
if self._pending_writes > 0:
await self.commit()
except asyncio.CancelledError:
raise
except Exception:
logger.exception("Error in database flush timer.")
+1 -4
View File
@@ -92,10 +92,7 @@ async def _ingest_sentence(
start_words.append((channel_id, user_id, words[i])) start_words.append((channel_id, user_id, words[i]))
transitions.append((channel_id, user_id, words[i], words[i + 1])) transitions.append((channel_id, user_id, words[i], words[i + 1]))
if start_words: await db.add_markov_data(start_words, transitions)
await db.add_start_words_batch(start_words)
if transitions:
await db.add_transitions_batch(transitions)
async def generate( async def generate(
+42 -90
View File
@@ -15,16 +15,14 @@
"""Unit tests for the Database class. """Unit tests for the Database class.
Tests cover Database.connect (pragmas, schema), markov start word and Tests cover Database.connect (pragmas, schema), markov start word and
transition CRUD, image storage, flag CRUD, channel ingestion tracking, transition CRUD, image storage, flag CRUD, and channel ingestion tracking.
and the auto-commit write-batching mechanism.
""" """
import asyncio
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
import pytest import pytest
from crabstero.database import AUTO_COMMIT_WRITE_THRESHOLD, Database from crabstero.database import Database
if TYPE_CHECKING: if TYPE_CHECKING:
from pathlib import Path from pathlib import Path
@@ -33,13 +31,6 @@ if TYPE_CHECKING:
class TestConnect: class TestConnect:
"""Database.connect creates a configured SQLite database.""" """Database.connect creates a configured SQLite database."""
async def test_wal_mode(self, db: Database) -> None:
"""Journal mode is set to WAL."""
async with db._connection.execute("PRAGMA journal_mode") as cursor:
row = await cursor.fetchone()
assert row is not None
assert row[0] == "wal"
async def test_synchronous_normal(self, db: Database) -> None: async def test_synchronous_normal(self, db: Database) -> None:
"""Synchronous mode is set to NORMAL.""" """Synchronous mode is set to NORMAL."""
async with db._connection.execute("PRAGMA synchronous") as cursor: async with db._connection.execute("PRAGMA synchronous") as cursor:
@@ -62,42 +53,34 @@ class TestConnect:
] ]
class TestMarkovStartWords: class TestAddMarkovData:
"""Start word storage and random retrieval.""" """Markov start word and transition storage via add_markov_data."""
async def test_add_and_retrieve(self, db: Database) -> None: async def test_stores_start_word(self, db: Database) -> None:
"""Inserted start word can be retrieved by channel.""" """Inserted start word can be retrieved by channel."""
await db.add_start_words_batch([(1, 100, "Hello")]) await db.add_markov_data([(1, 100, "Hello")], [])
await db.commit()
result = await db.get_random_start_word(1) result = await db.get_random_start_word(1)
assert result == "Hello" assert result == "Hello"
async def test_returns_none_when_empty(self, db: Database) -> None: async def test_stores_transition(self, db: Database) -> None:
"""Returns None for a channel with no start words."""
result = await db.get_random_start_word(999)
assert result is None
async def test_empty_batch_stores_nothing(self, db: Database) -> None:
"""An empty batch does not store any start words."""
await db.add_start_words_batch([])
await db.commit()
assert await db.get_random_start_word(1) is None
class TestMarkovTransitions:
"""Word transition storage and random retrieval."""
async def test_add_and_retrieve(self, db: Database) -> None:
"""Inserted transition can be retrieved by channel and word.""" """Inserted transition can be retrieved by channel and word."""
await db.add_transitions_batch([(1, 100, "Hello", "world.")]) await db.add_markov_data([], [(1, 100, "Hello", "world.")])
await db.commit()
result = await db.get_random_next_word(1, "Hello") result = await db.get_random_next_word(1, "Hello")
assert result == "world." assert result == "world."
async def test_returns_none_when_empty(self, db: Database) -> None: async def test_stores_both_in_single_call(self, db: Database) -> None:
"""Returns None when no transitions exist for the word.""" """Start words and transitions are stored in a single call."""
result = await db.get_random_next_word(1, "nonexistent") await db.add_markov_data(
assert result is None [(1, 100, "Hello")],
[(1, 100, "Hello", "world.")],
)
assert await db.get_random_start_word(1) == "Hello"
assert await db.get_random_next_word(1, "Hello") == "world."
async def test_empty_lists_store_nothing(self, db: Database) -> None:
"""Empty lists do not store any data."""
await db.add_markov_data([], [])
assert await db.get_random_start_word(1) is None
@pytest.mark.parametrize( @pytest.mark.parametrize(
"completing_word", "completing_word",
@@ -105,36 +88,38 @@ class TestMarkovTransitions:
pytest.param("world.", id="period"), pytest.param("world.", id="period"),
pytest.param("world!", id="exclamation"), pytest.param("world!", id="exclamation"),
pytest.param("world?", id="question"), pytest.param("world?", id="question"),
pytest.param("world§", id="section-sign"), pytest.param("world\u00a7", id="section-sign"),
], ],
) )
async def test_completing_word_filters_punctuation( async def test_completing_word_filters_punctuation(
self, db: Database, completing_word: str self, db: Database, completing_word: str
) -> None: ) -> None:
"""get_random_completing_next_word only returns sentence-ending words.""" """get_random_completing_next_word only returns sentence-ending words."""
await db.add_transitions_batch( await db.add_markov_data(
[],
[ [
(1, 100, "Hello", "beautiful"), (1, 100, "Hello", "beautiful"),
(1, 100, "Hello", completing_word), (1, 100, "Hello", completing_word),
] ],
) )
await db.commit()
for _ in range(100): for _ in range(100):
result = await db.get_random_completing_next_word(1, "Hello") result = await db.get_random_completing_next_word(1, "Hello")
assert result == completing_word assert result == completing_word
async def test_completing_returns_none_without_match(self, db: Database) -> None: async def test_completing_returns_none_without_match(self, db: Database) -> None:
"""Returns None when no transitions end with sentence punctuation.""" """Returns None when no transitions end with sentence punctuation."""
await db.add_transitions_batch([(1, 100, "Hello", "beautiful")]) await db.add_markov_data([], [(1, 100, "Hello", "beautiful")])
await db.commit()
result = await db.get_random_completing_next_word(1, "Hello") result = await db.get_random_completing_next_word(1, "Hello")
assert result is None assert result is None
async def test_empty_batch_stores_nothing(self, db: Database) -> None:
"""An empty batch does not store any transitions.""" class TestMarkovReadMethods:
await db.add_transitions_batch([]) """Read behavior when no data has been stored."""
await db.commit()
assert await db.get_random_next_word(1, "Hello") is None async def test_returns_none_when_empty(self, db: Database) -> None:
"""Returns None for channels with no data."""
assert await db.get_random_start_word(999) is None
assert await db.get_random_next_word(1, "nonexistent") is None
class TestImages: class TestImages:
@@ -143,7 +128,6 @@ class TestImages:
async def test_add_and_retrieve(self, db: Database) -> None: async def test_add_and_retrieve(self, db: Database) -> None:
"""Inserted image URL can be retrieved by channel.""" """Inserted image URL can be retrieved by channel."""
await db.add_image(1, 100, "https://example.com/cat.png") await db.add_image(1, 100, "https://example.com/cat.png")
await db.commit()
result = await db.get_random_image(1) result = await db.get_random_image(1)
assert result == "https://example.com/cat.png" assert result == "https://example.com/cat.png"
@@ -159,7 +143,6 @@ class TestFlags:
async def test_set_and_check(self, db: Database) -> None: async def test_set_and_check(self, db: Database) -> None:
"""A set flag is reported as set.""" """A set flag is reported as set."""
await db.set_flag("channel", "123", "noReply") await db.set_flag("channel", "123", "noReply")
await db.commit()
assert await db.is_flag_set("channel", "123", "noReply") is True assert await db.is_flag_set("channel", "123", "noReply") is True
async def test_unset_flag_is_false(self, db: Database) -> None: async def test_unset_flag_is_false(self, db: Database) -> None:
@@ -169,16 +152,13 @@ class TestFlags:
async def test_clear_flag(self, db: Database) -> None: async def test_clear_flag(self, db: Database) -> None:
"""A cleared flag is no longer reported as set.""" """A cleared flag is no longer reported as set."""
await db.set_flag("channel", "123", "noReply") await db.set_flag("channel", "123", "noReply")
await db.commit()
await db.clear_flag("channel", "123", "noReply") await db.clear_flag("channel", "123", "noReply")
await db.commit()
assert await db.is_flag_set("channel", "123", "noReply") is False assert await db.is_flag_set("channel", "123", "noReply") is False
async def test_set_idempotent(self, db: Database) -> None: async def test_set_idempotent(self, db: Database) -> None:
"""Setting the same flag twice does not raise.""" """Setting the same flag twice does not raise."""
await db.set_flag("channel", "123", "noReply") await db.set_flag("channel", "123", "noReply")
await db.set_flag("channel", "123", "noReply") await db.set_flag("channel", "123", "noReply")
await db.commit()
assert await db.is_flag_set("channel", "123", "noReply") is True assert await db.is_flag_set("channel", "123", "noReply") is True
@@ -188,7 +168,6 @@ class TestChannelIngestion:
async def test_mark_and_check(self, db: Database) -> None: async def test_mark_and_check(self, db: Database) -> None:
"""A marked channel is reported as ingested.""" """A marked channel is reported as ingested."""
await db.mark_channel_ingested(42) await db.mark_channel_ingested(42)
await db.commit()
assert await db.is_channel_ingested(42) is True assert await db.is_channel_ingested(42) is True
async def test_not_ingested_by_default(self, db: Database) -> None: async def test_not_ingested_by_default(self, db: Database) -> None:
@@ -199,47 +178,20 @@ class TestChannelIngestion:
"""Marking the same channel twice does not raise.""" """Marking the same channel twice does not raise."""
await db.mark_channel_ingested(42) await db.mark_channel_ingested(42)
await db.mark_channel_ingested(42) await db.mark_channel_ingested(42)
await db.commit()
assert await db.is_channel_ingested(42) is True assert await db.is_channel_ingested(42) is True
class TestAutoCommit: class TestWriteDurability:
"""Write batching and automatic commit behavior.""" """Writes persist across close and reopen."""
async def test_commits_after_threshold(self, db: Database) -> None: async def test_markov_data_survives_reopen(self, tmp_path: Path) -> None:
"""Pending writes reset to zero after reaching the write threshold.""" """Data written via add_markov_data is durable after close/reopen."""
for i in range(AUTO_COMMIT_WRITE_THRESHOLD): db_path = str(tmp_path / "durability.db")
await db.add_start_words_batch([(1, 100, f"word{i}")])
assert db._pending_writes == 0
assert await db.get_random_start_word(1) is not None
async def test_close_commits_pending(self, tmp_path: Path) -> None:
"""close() commits any pending writes before closing."""
db_path = str(tmp_path / "close_test.db")
db = await Database.connect(db_path) db = await Database.connect(db_path)
await db.add_start_words_batch([(1, 100, "hello")]) await db.add_markov_data([(1, 100, "Hello")], [(1, 100, "Hello", "world.")])
assert db._pending_writes > 0
await db.close() await db.close()
# Reopen and verify data persisted.
db2 = await Database.connect(db_path) db2 = await Database.connect(db_path)
result = await db2.get_random_start_word(1) assert await db2.get_random_start_word(1) == "Hello"
assert await db2.get_random_next_word(1, "Hello") == "world."
await db2.close() await db2.close()
assert result == "hello"
async def test_close_without_pending_writes(self, tmp_path: Path) -> None:
"""close() succeeds when there are no pending writes."""
db = await Database.connect(str(tmp_path / "clean_close.db"))
await db.close()
async def test_flush_timer_commits(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Background flush timer commits pending writes after timeout."""
monkeypatch.setattr("crabstero.database.AUTO_COMMIT_TIMEOUT_SECONDS", 0.05)
db = await Database.connect(str(tmp_path / "timer_test.db"))
await db.add_start_words_batch([(1, 100, "hello")])
assert db._pending_writes > 0
await asyncio.sleep(0.1)
assert db._pending_writes == 0
await db.close()
-3
View File
@@ -103,7 +103,6 @@ class TestSetClearCheck:
async def test_set_then_check(self, db: Database, flag: Flag) -> None: async def test_set_then_check(self, db: Database, flag: Flag) -> None:
"""A set flag is reported as set.""" """A set flag is reported as set."""
await set_flag(db, 1, EntityType.CHANNEL, flag) await set_flag(db, 1, EntityType.CHANNEL, flag)
await db.commit()
assert await is_flag_set(db, 1, EntityType.CHANNEL, flag) is True assert await is_flag_set(db, 1, EntityType.CHANNEL, flag) is True
async def test_unset_returns_false(self, db: Database) -> None: async def test_unset_returns_false(self, db: Database) -> None:
@@ -113,7 +112,5 @@ class TestSetClearCheck:
async def test_clear_removes_flag(self, db: Database) -> None: async def test_clear_removes_flag(self, db: Database) -> None:
"""A cleared flag is no longer reported as set.""" """A cleared flag is no longer reported as set."""
await set_flag(db, 1, EntityType.CHANNEL, Flag.NO_REPLY) await set_flag(db, 1, EntityType.CHANNEL, Flag.NO_REPLY)
await db.commit()
await clear_flag(db, 1, EntityType.CHANNEL, Flag.NO_REPLY) await clear_flag(db, 1, EntityType.CHANNEL, Flag.NO_REPLY)
await db.commit()
assert await is_flag_set(db, 1, EntityType.CHANNEL, Flag.NO_REPLY) is False assert await is_flag_set(db, 1, EntityType.CHANNEL, Flag.NO_REPLY) is False
+33 -35
View File
@@ -61,42 +61,42 @@ class TestIngestSentence:
async def test_stores_start_word(self, db: Database) -> None: async def test_stores_start_word(self, db: Database) -> None:
"""First word of the sentence is stored as a start word.""" """First word of the sentence is stored as a start word."""
await _ingest_sentence(db, 1, 100, "Hello world.") await _ingest_sentence(db, 1, 100, "Hello world.")
await db.commit()
result = await db.get_random_start_word(1) result = await db.get_random_start_word(1)
assert result == "Hello" assert result == "Hello"
async def test_stores_transitions(self, db: Database) -> None: async def test_stores_transitions(self, db: Database) -> None:
"""Adjacent words create transitions.""" """Adjacent words create transitions."""
await _ingest_sentence(db, 1, 100, "Hello world.") await _ingest_sentence(db, 1, 100, "Hello world.")
await db.commit()
result = await db.get_random_next_word(1, "Hello") result = await db.get_random_next_word(1, "Hello")
assert result == "world." assert result == "world."
async def test_appends_sentence_end_if_missing(self, db: Database) -> None: async def test_appends_sentence_end_if_missing(self, db: Database) -> None:
"""Unpunctuated sentence gets the default sentence-end marker.""" """Unpunctuated sentence gets the default sentence-end marker."""
await _ingest_sentence(db, 1, 100, "Hello world") await _ingest_sentence(db, 1, 100, "Hello world")
await db.commit()
result = await db.get_random_next_word(1, "Hello") result = await db.get_random_next_word(1, "Hello")
assert result == f"world{DEFAULT_SENTENCE_END}" assert result == f"world{DEFAULT_SENTENCE_END}"
async def test_preserves_existing_punctuation(self, db: Database) -> None: async def test_preserves_existing_punctuation(self, db: Database) -> None:
"""Already-punctuated sentence keeps its terminator.""" """Already-punctuated sentence keeps its terminator."""
await _ingest_sentence(db, 1, 100, "Hello world!") await _ingest_sentence(db, 1, 100, "Hello world!")
await db.commit()
result = await db.get_random_next_word(1, "Hello") result = await db.get_random_next_word(1, "Hello")
assert result == "world!" assert result == "world!"
async def test_single_word_stores_nothing(self, db: Database) -> None: async def test_single_word_stores_nothing(self, db: Database) -> None:
"""A single-word sentence produces no start words or transitions.""" """A single-word sentence produces no start words or transitions."""
await _ingest_sentence(db, 1, 100, "Hello.") await _ingest_sentence(db, 1, 100, "Hello.")
await db.commit()
assert await db.get_random_start_word(1) is None assert await db.get_random_start_word(1) is None
assert await db.get_random_next_word(1, "Hello.") is None assert await db.get_random_next_word(1, "Hello.") is None
async def test_stores_all_transitions(self, db: Database) -> None: async def test_stores_all_transitions(self, db: Database) -> None:
"""All adjacent word pairs create transitions.""" """All adjacent word pairs create transitions."""
await _ingest_sentence(db, 1, 100, "A B C.") await _ingest_sentence(db, 1, 100, "A B C.")
await db.commit()
assert await db.get_random_next_word(1, "A") == "B" assert await db.get_random_next_word(1, "A") == "B"
assert await db.get_random_next_word(1, "B") == "C." assert await db.get_random_next_word(1, "B") == "C."
@@ -107,13 +107,13 @@ class TestIngest:
async def test_single_sentence(self, db: Database) -> None: async def test_single_sentence(self, db: Database) -> None:
"""A single sentence paragraph is ingested.""" """A single sentence paragraph is ingested."""
await ingest(db, 1, 100, "Hello world.") await ingest(db, 1, 100, "Hello world.")
await db.commit()
assert await db.get_random_start_word(1) == "Hello" assert await db.get_random_start_word(1) == "Hello"
async def test_multiple_sentences(self, db: Database) -> None: async def test_multiple_sentences(self, db: Database) -> None:
"""Multiple sentences are split and ingested individually.""" """Multiple sentences are split and ingested individually."""
await ingest(db, 1, 100, "Hello world. Goodbye world!") await ingest(db, 1, 100, "Hello world. Goodbye world!")
await db.commit()
# Both "Hello" and "Goodbye" should appear as start words. # Both "Hello" and "Goodbye" should appear as start words.
async with db._connection.execute( async with db._connection.execute(
"SELECT DISTINCT word FROM markov_start_words WHERE channel_id = ?", "SELECT DISTINCT word FROM markov_start_words WHERE channel_id = ?",
@@ -125,14 +125,14 @@ class TestIngest:
async def test_normalizes_whitespace(self, db: Database) -> None: async def test_normalizes_whitespace(self, db: Database) -> None:
"""Extra spaces and newlines are collapsed.""" """Extra spaces and newlines are collapsed."""
await ingest(db, 1, 100, "Hello world.\nGoodbye world!") await ingest(db, 1, 100, "Hello world.\nGoodbye world!")
await db.commit()
result = await db.get_random_next_word(1, "Hello") result = await db.get_random_next_word(1, "Hello")
assert result == "world." assert result == "world."
async def test_appends_default_end(self, db: Database) -> None: async def test_appends_default_end(self, db: Database) -> None:
"""Unpunctuated paragraph gets the default sentence-end marker.""" """Unpunctuated paragraph gets the default sentence-end marker."""
await ingest(db, 1, 100, "Hello world") await ingest(db, 1, 100, "Hello world")
await db.commit()
result = await db.get_random_next_word(1, "Hello") result = await db.get_random_next_word(1, "Hello")
assert result == f"world{DEFAULT_SENTENCE_END}" assert result == f"world{DEFAULT_SENTENCE_END}"
@@ -148,13 +148,13 @@ class TestIngest:
) -> None: ) -> None:
"""Empty or whitespace-only input does not store any data.""" """Empty or whitespace-only input does not store any data."""
await ingest(db, 1, 100, paragraph) await ingest(db, 1, 100, paragraph)
await db.commit()
assert await db.get_random_start_word(1) is None assert await db.get_random_start_word(1) is None
async def test_splits_on_punctuation_followed_by_space(self, db: Database) -> None: async def test_splits_on_punctuation_followed_by_space(self, db: Database) -> None:
"""Punctuation followed by a space splits into separate sentences.""" """Punctuation followed by a space splits into separate sentences."""
await ingest(db, 1, 100, "Dr. Smith likes cats") await ingest(db, 1, 100, "Dr. Smith likes cats")
await db.commit()
# "Dr." splits off as a single-word sentence (stores nothing). # "Dr." splits off as a single-word sentence (stores nothing).
# "Smith likes cats" becomes a sentence with "Smith" as start word. # "Smith likes cats" becomes a sentence with "Smith" as start word.
assert await db.get_random_start_word(1) == "Smith" assert await db.get_random_start_word(1) == "Smith"
@@ -162,7 +162,7 @@ class TestIngest:
async def test_no_split_without_space_after_punctuation(self, db: Database) -> None: async def test_no_split_without_space_after_punctuation(self, db: Database) -> None:
"""Punctuation not followed by a space keeps words together.""" """Punctuation not followed by a space keeps words together."""
await ingest(db, 1, 100, "Hello.World is here") await ingest(db, 1, 100, "Hello.World is here")
await db.commit()
assert await db.get_random_start_word(1) == "Hello.World" assert await db.get_random_start_word(1) == "Hello.World"
@@ -177,14 +177,14 @@ class TestGenerate:
async def test_generates_from_ingested_data(self, db: Database) -> None: async def test_generates_from_ingested_data(self, db: Database) -> None:
"""Generated text uses words from ingested data.""" """Generated text uses words from ingested data."""
await ingest(db, 1, 100, "The quick brown fox.") await ingest(db, 1, 100, "The quick brown fox.")
await db.commit()
result = await generate(db, 1) result = await generate(db, 1)
assert result == "The quick brown fox." assert result == "The quick brown fox."
async def test_strips_section_sign(self, db: Database) -> None: async def test_strips_section_sign(self, db: Database) -> None:
"""The internal section sign marker never appears in output.""" """The internal section sign marker never appears in output."""
await ingest(db, 1, 100, "Hello world") await ingest(db, 1, 100, "Hello world")
await db.commit()
result = await generate(db, 1) result = await generate(db, 1)
assert result == "Hello world" assert result == "Hello world"
@@ -193,7 +193,7 @@ class TestGenerate:
words = [f"w{i}" for i in range(100)] words = [f"w{i}" for i in range(100)]
text = " ".join(words) text = " ".join(words)
await ingest(db, 1, 100, text) await ingest(db, 1, 100, text)
await db.commit()
result = await generate(db, 1, soft_limit=10, hard_limit=49) result = await generate(db, 1, soft_limit=10, hard_limit=49)
assert result == "w0 w1 w2 w3 w4 w5 w6 w7 w8 w9 w10 w11 w12 w13 w14" assert result == "w0 w1 w2 w3 w4 w5 w6 w7 w8 w9 w10 w11 w12 w13 w14"
@@ -204,8 +204,8 @@ class TestGenerate:
# the loop falls back to get_random_next_word for each. D has both a # the loop falls back to get_random_next_word for each. D has both a
# continuing ("E") and completing ("end.") transition, so # continuing ("E") and completing ("end.") transition, so
# get_random_completing_next_word deterministically picks "end.". # get_random_completing_next_word deterministically picks "end.".
await db.add_start_words_batch([(1, 100, "A")]) await db.add_markov_data(
await db.add_transitions_batch( [(1, 100, "A")],
[ [
(1, 100, "A", "B"), (1, 100, "A", "B"),
(1, 100, "B", "C"), (1, 100, "B", "C"),
@@ -215,24 +215,24 @@ class TestGenerate:
(1, 100, "D", "G"), (1, 100, "D", "G"),
(1, 100, "D", "H"), (1, 100, "D", "H"),
(1, 100, "D", "end."), (1, 100, "D", "end."),
] ],
) )
await db.commit()
for _ in range(100): for _ in range(100):
result = await generate(db, 1, soft_limit=1, hard_limit=1000) result = await generate(db, 1, soft_limit=1, hard_limit=1000)
assert result == "A B C D end." assert result == "A B C D end."
async def test_start_word_already_ends_sentence(self, db: Database) -> None: async def test_start_word_already_ends_sentence(self, db: Database) -> None:
"""Generation stops immediately when the start word is sentence-ending.""" """Generation stops immediately when the start word is sentence-ending."""
await db.add_start_words_batch([(1, 100, "Yes.")]) await db.add_markov_data([(1, 100, "Yes.")], [])
await db.commit()
result = await generate(db, 1) result = await generate(db, 1)
assert result == "Yes." assert result == "Yes."
async def test_chain_dead_end(self, db: Database) -> None: async def test_chain_dead_end(self, db: Database) -> None:
"""Generation stops when no next word exists (dead-end chain).""" """Generation stops when no next word exists (dead-end chain)."""
await db.add_start_words_batch([(1, 100, "Hello")]) await db.add_markov_data([(1, 100, "Hello")], [])
await db.commit()
# "Hello" has no transitions, so the loop breaks immediately. # "Hello" has no transitions, so the loop breaks immediately.
result = await generate(db, 1) result = await generate(db, 1)
assert result == "Hello" assert result == "Hello"
@@ -240,14 +240,14 @@ class TestGenerate:
async def test_hard_limit_strips_section_sign(self, db: Database) -> None: async def test_hard_limit_strips_section_sign(self, db: Database) -> None:
"""Section sign at the truncation boundary is stripped.""" """Section sign at the truncation boundary is stripped."""
# Build a chain: "A" -> "B" -> "C§". # Build a chain: "A" -> "B" -> "C§".
await db.add_start_words_batch([(1, 100, "A")]) await db.add_markov_data(
await db.add_transitions_batch( [(1, 100, "A")],
[ [
(1, 100, "A", "B"), (1, 100, "A", "B"),
(1, 100, "B", f"C{DEFAULT_SENTENCE_END}"), (1, 100, "B", f"C{DEFAULT_SENTENCE_END}"),
] ],
) )
await db.commit()
# hard_limit=5 truncates "A B C§" (length 6) to "A B C". # hard_limit=5 truncates "A B C§" (length 6) to "A B C".
result = await generate(db, 1, soft_limit=100, hard_limit=5) result = await generate(db, 1, soft_limit=100, hard_limit=5)
assert result == "A B C" assert result == "A B C"
@@ -258,17 +258,15 @@ class TestGenerate:
"""After soft_limit, falls back when no completing word exists.""" """After soft_limit, falls back when no completing word exists."""
# "A" -> "B" (no completing transition). Past soft_limit, # "A" -> "B" (no completing transition). Past soft_limit,
# get_random_completing_next_word returns None, falls back to "B". # get_random_completing_next_word returns None, falls back to "B".
await db.add_start_words_batch([(1, 100, "A")]) await db.add_markov_data([(1, 100, "A")], [(1, 100, "A", "B")])
await db.add_transitions_batch([(1, 100, "A", "B")])
await db.commit()
result = await generate(db, 1, soft_limit=1, hard_limit=1000) result = await generate(db, 1, soft_limit=1, hard_limit=1000)
assert result == "A B" assert result == "A B"
async def test_hard_limit_truncates_mid_word(self, db: Database) -> None: async def test_hard_limit_truncates_mid_word(self, db: Database) -> None:
"""Hard limit slices output even when it falls inside a word.""" """Hard limit slices output even when it falls inside a word."""
await db.add_start_words_batch([(1, 100, "AB")]) await db.add_markov_data([(1, 100, "AB")], [(1, 100, "AB", "CDEF")])
await db.add_transitions_batch([(1, 100, "AB", "CDEF")])
await db.commit()
# "AB CDEF" is 7 chars; hard_limit=5 truncates to "AB CD". # "AB CDEF" is 7 chars; hard_limit=5 truncates to "AB CD".
result = await generate(db, 1, soft_limit=100, hard_limit=5) result = await generate(db, 1, soft_limit=100, hard_limit=5)
assert result == "AB CD" assert result == "AB CD"
@@ -277,7 +275,7 @@ class TestGenerate:
"""Data ingested into one channel does not leak into another.""" """Data ingested into one channel does not leak into another."""
await ingest(db, 1, 100, "Channel one data.") await ingest(db, 1, 100, "Channel one data.")
await ingest(db, 2, 100, "Channel two data.") await ingest(db, 2, 100, "Channel two data.")
await db.commit()
# Channel 3 has no data; should get the fallback. # Channel 3 has no data; should get the fallback.
result = await generate(db, 3) result = await generate(db, 3)
assert result == "Hello world!" assert result == "Hello world!"