Simplified database layer and consolidated Markov writes into a single transaction.
This commit is contained in:
+28
-82
@@ -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.
|
||||||
"""
|
"""
|
||||||
|
if start_words:
|
||||||
await self._connection.executemany(
|
await self._connection.executemany(
|
||||||
"INSERT INTO markov_start_words"
|
"INSERT INTO markov_start_words"
|
||||||
" (channel_id, user_id, word)"
|
" (channel_id, user_id, word)"
|
||||||
" VALUES (?, ?, ?)",
|
" VALUES (?, ?, ?)",
|
||||||
rows,
|
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
@@ -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
@@ -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()
|
|
||||||
|
|||||||
@@ -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
@@ -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!"
|
||||||
|
|||||||
Reference in New Issue
Block a user