1288 lines
42 KiB
Python
1288 lines
42 KiB
Python
"""Tests for the database module."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
if TYPE_CHECKING:
|
|
import sqlite3
|
|
|
|
from vibe_bot.database import ChatDatabase
|
|
|
|
|
|
def _recent_messages(
|
|
db_path: str,
|
|
limit: int,
|
|
) -> list[tuple[str, str, str, datetime]]:
|
|
"""Read the newest rows straight from the database (test-side helper)."""
|
|
from vibe_bot.db.connection import connect
|
|
|
|
conn = connect(db_path)
|
|
try:
|
|
rows = conn.execute(
|
|
"SELECT message_id, username, content, timestamp "
|
|
"FROM chat_messages ORDER BY timestamp DESC LIMIT ?",
|
|
(limit,),
|
|
).fetchall()
|
|
finally:
|
|
conn.close()
|
|
return [
|
|
(str(row[0]), str(row[1]), str(row[2]), cast(datetime, row[3])) for row in rows
|
|
]
|
|
|
|
|
|
def test_vector_to_bytes(chat_db: ChatDatabase) -> None:
|
|
"""Test converting a vector to bytes and back."""
|
|
vector: list[float] = [0.1, 0.2, 0.3, 0.4]
|
|
blob = chat_db._vector_to_bytes(vector)
|
|
assert isinstance(blob, bytes)
|
|
assert len(blob) == len(vector) * 4 # float32 = 4 bytes
|
|
|
|
reconstructed = chat_db._bytes_to_vector(blob)
|
|
assert np.allclose(reconstructed, np.array(vector, dtype=np.float32))
|
|
|
|
|
|
def test_bytes_to_vector(chat_db: ChatDatabase) -> None:
|
|
"""Test converting bytes back to a numpy vector."""
|
|
original = np.array([1.0, 2.0, 3.0], dtype=np.float32)
|
|
blob = original.tobytes()
|
|
result = chat_db._bytes_to_vector(blob)
|
|
assert np.array_equal(result, original)
|
|
|
|
|
|
def test_calculate_similarity_self(chat_db: ChatDatabase) -> None:
|
|
"""Test cosine similarity of a vector with itself is 1.0."""
|
|
vec = np.array([1.0, 2.0, 3.0], dtype=np.float32)
|
|
similarity = chat_db._calculate_similarity(vec, vec)
|
|
assert similarity == pytest.approx(1.0, abs=1e-6)
|
|
|
|
|
|
def test_calculate_similarity_orthogonal(chat_db: ChatDatabase) -> None:
|
|
"""Test cosine similarity of orthogonal vectors is 0."""
|
|
vec1 = np.array([1.0, 0.0], dtype=np.float32)
|
|
vec2 = np.array([0.0, 1.0], dtype=np.float32)
|
|
similarity = chat_db._calculate_similarity(vec1, vec2)
|
|
assert similarity == pytest.approx(0.0, abs=1e-6)
|
|
|
|
|
|
def test_calculate_similarity_negative(chat_db: ChatDatabase) -> None:
|
|
"""Test cosine similarity of opposite vectors is -1."""
|
|
vec1 = np.array([1.0, 0.0], dtype=np.float32)
|
|
vec2 = np.array([-1.0, 0.0], dtype=np.float32)
|
|
similarity = chat_db._calculate_similarity(vec1, vec2)
|
|
assert similarity == pytest.approx(-1.0, abs=1e-6)
|
|
|
|
|
|
def test_calculate_similarity_zero_norm(chat_db: ChatDatabase) -> None:
|
|
"""A zero vector has no direction, so its similarity is 0."""
|
|
zero = np.zeros(3, dtype=np.float32)
|
|
other = np.array([1.0, 0.0, 0.0], dtype=np.float32)
|
|
assert chat_db._calculate_similarity(zero, other) == 0.0
|
|
|
|
|
|
def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
|
|
"""Test adding a message to the database."""
|
|
result = chat_db.add_message(
|
|
message_id="msg-1",
|
|
user_id="user-1",
|
|
username="testuser",
|
|
content="Hello world",
|
|
channel_id="chan-1",
|
|
guild_id="guild-1",
|
|
)
|
|
assert result is True
|
|
|
|
messages = _recent_messages(chat_db.db_path, 10)
|
|
assert len(messages) == 1
|
|
assert messages[0][0] == "msg-1"
|
|
assert messages[0][1] == "testuser"
|
|
assert messages[0][2] == "Hello world"
|
|
|
|
|
|
def test_add_message_no_embedding(chat_db: ChatDatabase) -> None:
|
|
"""Test adding a message when embedding generation fails."""
|
|
with patch("vibe_bot.llm_client.embedding", return_value=None):
|
|
result = chat_db.add_message(
|
|
message_id="msg-no-embed",
|
|
user_id="user-1",
|
|
username="testuser",
|
|
content="No embedding message",
|
|
channel_id="chan-1",
|
|
guild_id="guild-1",
|
|
)
|
|
assert result is True
|
|
|
|
|
|
def test_add_message_duplicate(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""Test adding a duplicate message replaces the old one."""
|
|
chat_db.add_message(
|
|
message_id="msg-dup",
|
|
user_id="user-1",
|
|
username="testuser",
|
|
content="First content",
|
|
)
|
|
chat_db.add_message(
|
|
message_id="msg-dup",
|
|
user_id="user-1",
|
|
username="testuser",
|
|
content="Second content",
|
|
)
|
|
|
|
messages = _recent_messages(chat_db.db_path, 10)
|
|
assert len(messages) == 1
|
|
assert messages[0][2] == "Second content"
|
|
|
|
|
|
def test_add_message_failure(chat_db: ChatDatabase) -> None:
|
|
"""Test that add_message returns False on database error."""
|
|
with patch("vibe_bot.db.messages.connect", return_value=_broken_connection()):
|
|
result = chat_db.add_message(
|
|
message_id="msg-fail",
|
|
user_id="user-1",
|
|
username="testuser",
|
|
content="Should fail",
|
|
)
|
|
assert result is False
|
|
|
|
|
|
def _embedding_row_count(db_path: str) -> int:
|
|
"""Count the rows stored in message_embeddings."""
|
|
import sqlite3
|
|
|
|
conn = sqlite3.connect(db_path)
|
|
count = conn.execute("SELECT COUNT(*) FROM message_embeddings").fetchone()[0]
|
|
conn.close()
|
|
return int(count)
|
|
|
|
|
|
def test_add_message_embed_false_skips_embedding(chat_db: ChatDatabase) -> None:
|
|
"""embed=False neither calls the embedding API nor stores a row."""
|
|
with patch("vibe_bot.llm_client.embedding") as mock_embedding:
|
|
result = chat_db.add_message(
|
|
message_id="msg-assist",
|
|
user_id="bot-1",
|
|
username="some-bot",
|
|
content="assistant reply",
|
|
role="assistant",
|
|
embed=False,
|
|
)
|
|
assert result is True
|
|
mock_embedding.assert_not_called()
|
|
assert _embedding_row_count(chat_db.db_path) == 0
|
|
|
|
|
|
def test_add_message_stores_embedding_by_default(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""The default embed=True still stores exactly one embedding row."""
|
|
assert chat_db.add_message(
|
|
message_id="msg-embed",
|
|
user_id="user-1",
|
|
username="testuser",
|
|
content="default embed",
|
|
)
|
|
assert _embedding_row_count(chat_db.db_path) == 1
|
|
|
|
|
|
def test_add_message_stores_embedding_norm(chat_db: ChatDatabase) -> None:
|
|
"""add_message stores the L2 norm of the float32-encoded embedding."""
|
|
import sqlite3
|
|
|
|
vector: list[float] = [0.6, 0.8]
|
|
with patch("vibe_bot.llm_client.embedding", return_value=vector):
|
|
assert chat_db.add_message(
|
|
message_id="norm-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="normed message",
|
|
)
|
|
|
|
conn = sqlite3.connect(chat_db.db_path)
|
|
blob, norm = conn.execute(
|
|
"SELECT embedding, norm FROM message_embeddings WHERE message_id = 'norm-1'"
|
|
).fetchone()
|
|
conn.close()
|
|
|
|
stored = np.frombuffer(blob, dtype=np.float32)
|
|
assert float(norm) == pytest.approx(float(np.linalg.norm(stored)), abs=1e-6)
|
|
|
|
|
|
def test_cleanup_old_messages_no_orphaned_embeddings(chat_db: ChatDatabase) -> None:
|
|
"""Deleting the oldest rows must delete their embeddings, not the next ones.
|
|
|
|
Regression test for the bug where the embedding cleanup re-queried
|
|
``chat_messages`` *after* the message delete, so it stripped embeddings
|
|
from the next-oldest live rows while leaving the deleted rows' embeddings
|
|
behind as orphans.
|
|
"""
|
|
import sqlite3
|
|
|
|
with patch("vibe_bot.db.messages.MAX_HISTORY_MESSAGES", 5):
|
|
# Seed 7 rows with distinct ascending timestamps and an embedding for
|
|
# each, so the "oldest" ordering is deterministic.
|
|
conn = sqlite3.connect(chat_db.db_path)
|
|
cursor = conn.cursor()
|
|
for i in range(1, 8):
|
|
ts = f"2024-01-0{i} 00:00:00"
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages "
|
|
"(message_id, user_id, username, content, bot_name, timestamp) "
|
|
"VALUES (?, ?, ?, ?, ?, ?)",
|
|
(f"m{i}", "u1", "alice", f"content {i}", "bot", ts),
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO message_embeddings (message_id, embedding) "
|
|
"VALUES (?, ?)",
|
|
(f"m{i}", chat_db._vector_to_bytes([0.1, 0.2, 0.3])),
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
# The 8th insert pushes the count to 8 (> 5), so add_message's
|
|
# cleanup must delete exactly the 3 oldest rows and their embeddings.
|
|
assert chat_db.add_message(
|
|
message_id="m8",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="content 8",
|
|
)
|
|
|
|
conn = sqlite3.connect(chat_db.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT message_id FROM chat_messages")
|
|
live = {row[0] for row in cursor.fetchall()}
|
|
cursor.execute("SELECT message_id FROM message_embeddings")
|
|
embedded = {row[0] for row in cursor.fetchall()}
|
|
conn.close()
|
|
|
|
# The three oldest rows are gone; the rest (incl. the new one) remain.
|
|
assert live == {"m4", "m5", "m6", "m7", "m8"}
|
|
# No orphaned embeddings and no stripped survivors: the embedding set
|
|
# exactly matches the live message rows.
|
|
assert embedded == live
|
|
|
|
|
|
def test_role_scopes_history_and_search(chat_db: ChatDatabase) -> None:
|
|
"""get_user_history excludes responses; search matches only user rows."""
|
|
chat_db.add_message(
|
|
message_id="r-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="User asks about the weather",
|
|
role="user",
|
|
)
|
|
chat_db.add_message(
|
|
message_id="r-1_response",
|
|
user_id="bot",
|
|
username="some-bot",
|
|
content="Bot answers the weather",
|
|
role="assistant",
|
|
)
|
|
|
|
# get_user_history returns only the user row (paired with its response).
|
|
conversations = chat_db.get_user_history("u1")
|
|
assert len(conversations) == 1
|
|
assert conversations[0][0] == "User asks about the weather"
|
|
assert conversations[0][1] == "Bot answers the weather"
|
|
|
|
# search_similar_messages only considers user rows, never responses.
|
|
results = chat_db.search_similar_messages(
|
|
"User asks about the weather", top_k=5, min_similarity=0.0
|
|
)
|
|
assert len(results) == 1
|
|
assert results[0][0] == "User asks about the weather"
|
|
assert results[0][1] == "Bot answers the weather"
|
|
|
|
|
|
def test_role_migration_backfills_legacy_rows(
|
|
temp_db_path: str,
|
|
) -> None:
|
|
"""Legacy rows (no role column) are backfilled when ChatDatabase inits."""
|
|
import sqlite3
|
|
|
|
from vibe_bot.database import ChatDatabase
|
|
|
|
# Create the legacy schema (no role column) with a user row and its
|
|
# response row inserted directly.
|
|
conn = sqlite3.connect(temp_db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"CREATE TABLE chat_messages ("
|
|
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
|
|
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
|
|
"bot_name TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
|
|
"channel_id TEXT, guild_id TEXT)"
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages (message_id, user_id, username, content) "
|
|
"VALUES ('legacy-1', 'u1', 'alice', 'old question')"
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages (message_id, user_id, username, content) "
|
|
"VALUES ('legacy-1_response', 'bot', 'old-bot', 'old answer')"
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
# Initializing ChatDatabase should add and backfill the role column.
|
|
ChatDatabase(db_path=temp_db_path)
|
|
|
|
conn = sqlite3.connect(temp_db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT message_id, role FROM chat_messages ORDER BY message_id")
|
|
roles = {row[0]: row[1] for row in cursor.fetchall()}
|
|
conn.close()
|
|
|
|
assert roles["legacy-1"] == "user"
|
|
assert roles["legacy-1_response"] == "assistant"
|
|
|
|
|
|
def test_bot_name_migration_adds_column(
|
|
temp_db_path: str,
|
|
) -> None:
|
|
"""A pre-bot_name schema gets the column added when ChatDatabase inits."""
|
|
import sqlite3
|
|
|
|
from vibe_bot.database import ChatDatabase
|
|
|
|
conn = sqlite3.connect(temp_db_path)
|
|
conn.execute(
|
|
"CREATE TABLE chat_messages ("
|
|
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
|
|
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
|
|
"timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
|
|
"channel_id TEXT, guild_id TEXT, role TEXT)"
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
ChatDatabase(db_path=temp_db_path)
|
|
|
|
conn = sqlite3.connect(temp_db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("PRAGMA table_info(chat_messages)")
|
|
columns = {row[1] for row in cursor.fetchall()}
|
|
conn.close()
|
|
|
|
assert "bot_name" in columns
|
|
|
|
|
|
def _seed_pre_norm_db(db_path: str) -> list[tuple[str, str, list[float]]]:
|
|
"""Create a pre-norm-schema database (no norm column) with seeded rows."""
|
|
import sqlite3
|
|
|
|
rows: list[tuple[str, str, list[float]]] = [
|
|
("n-1", "ask about the sky", [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
|
|
("n-2", "ask about the sea", [0.7, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
|
|
("n-3", "ask about the sun", [0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
|
|
("n-4", "ask about the sand", [0.2, 0.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
|
|
("n-5", "ask about nothing", [0.0] * 8),
|
|
]
|
|
|
|
conn = sqlite3.connect(db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"CREATE TABLE chat_messages ("
|
|
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
|
|
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
|
|
"bot_name TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
|
|
"channel_id TEXT, guild_id TEXT, role TEXT)"
|
|
)
|
|
cursor.execute(
|
|
"CREATE TABLE message_embeddings ("
|
|
"message_id TEXT PRIMARY KEY, embedding BLOB,"
|
|
"FOREIGN KEY (message_id) REFERENCES chat_messages(message_id))"
|
|
)
|
|
for i, (message_id, content, vector) in enumerate(rows):
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages "
|
|
"(message_id, user_id, username, content, role, timestamp) "
|
|
"VALUES (?, 'u1', 'alice', ?, 'user', ?)",
|
|
(message_id, content, f"2024-01-0{i + 1} 00:00:00"),
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages "
|
|
"(message_id, user_id, username, content, role, timestamp) "
|
|
"VALUES (?, 'bot-1', 'some-bot', ?, 'assistant', ?)",
|
|
(
|
|
f"{message_id}_response",
|
|
f"response {i + 1}",
|
|
f"2024-01-0{i + 1} 00:00:01",
|
|
),
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO message_embeddings (message_id, embedding) VALUES (?, ?)",
|
|
(message_id, np.array(vector, dtype=np.float32).tobytes()),
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
return rows
|
|
|
|
|
|
def test_norm_migration_backfills_stored_norms(temp_db_path: str) -> None:
|
|
"""A pre-norm database gets the norm column added and backfilled on init."""
|
|
import sqlite3
|
|
|
|
from vibe_bot.database import ChatDatabase
|
|
|
|
rows = _seed_pre_norm_db(temp_db_path)
|
|
ChatDatabase(db_path=temp_db_path)
|
|
|
|
conn = sqlite3.connect(temp_db_path)
|
|
stored = dict(conn.execute("SELECT message_id, norm FROM message_embeddings"))
|
|
conn.close()
|
|
|
|
for message_id, _content, vector in rows:
|
|
expected = float(np.linalg.norm(np.array(vector, dtype=np.float32)))
|
|
assert stored[message_id] == pytest.approx(expected, abs=1e-5)
|
|
|
|
|
|
def test_recent_messages_desc_order(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""Newest-first ordering of stored messages."""
|
|
chat_db.add_message(
|
|
message_id="msg-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="First",
|
|
)
|
|
chat_db.add_message(
|
|
message_id="msg-2",
|
|
user_id="u2",
|
|
username="bob",
|
|
content="Second",
|
|
)
|
|
chat_db.add_message(
|
|
message_id="msg-3",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="Third",
|
|
)
|
|
|
|
messages = _recent_messages(chat_db.db_path, 2)
|
|
assert len(messages) == 2
|
|
assert messages[0][2] == "Third"
|
|
assert messages[1][2] == "Second"
|
|
|
|
|
|
def test_recent_messages_limit(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""The newest-rows query respects the limit."""
|
|
for i in range(5):
|
|
chat_db.add_message(
|
|
message_id=f"msg-{i}",
|
|
user_id="u1",
|
|
username="alice",
|
|
content=f"Message {i}",
|
|
)
|
|
|
|
messages = _recent_messages(chat_db.db_path, 3)
|
|
assert len(messages) == 3
|
|
|
|
|
|
def test_recent_messages_returns_datetime(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""The timestamp column comes back as a real datetime, not a string."""
|
|
chat_db.add_message(
|
|
message_id="dt-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="fresh message",
|
|
)
|
|
messages = _recent_messages(chat_db.db_path, 1)
|
|
assert isinstance(messages[0][3], datetime)
|
|
|
|
|
|
def test_clear_all_messages(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""Test clearing all messages."""
|
|
chat_db.add_message(
|
|
message_id="msg-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="Hello",
|
|
)
|
|
chat_db.add_message(
|
|
message_id="msg-2",
|
|
user_id="u2",
|
|
username="bob",
|
|
content="World",
|
|
)
|
|
|
|
chat_db.clear_all_messages()
|
|
|
|
messages = _recent_messages(chat_db.db_path, 10)
|
|
assert len(messages) == 0
|
|
|
|
|
|
def test_image_generation_time_estimate_empty(chat_db: ChatDatabase) -> None:
|
|
"""Test the estimate is None before any generations are recorded."""
|
|
assert chat_db.get_image_generation_time_estimate() is None
|
|
|
|
|
|
def test_record_and_get_image_generation_time_estimate(
|
|
chat_db: ChatDatabase,
|
|
) -> None:
|
|
"""Test recorded generation times produce an average estimate."""
|
|
assert chat_db.record_image_generation_time(10.0) is True
|
|
assert chat_db.record_image_generation_time(20.0) is True
|
|
|
|
estimate = chat_db.get_image_generation_time_estimate()
|
|
assert estimate == pytest.approx(15.0)
|
|
|
|
|
|
def test_image_generation_estimate_uses_recent_window(
|
|
chat_db: ChatDatabase,
|
|
) -> None:
|
|
"""Test only the most recent window of generations feeds the estimate."""
|
|
from vibe_bot.database import IMAGE_GEN_TIME_WINDOW
|
|
|
|
chat_db.record_image_generation_time(1000.0)
|
|
for _ in range(IMAGE_GEN_TIME_WINDOW):
|
|
chat_db.record_image_generation_time(10.0)
|
|
|
|
estimate = chat_db.get_image_generation_time_estimate()
|
|
assert estimate == pytest.approx(10.0)
|
|
|
|
|
|
def test_image_generation_times_capped(chat_db: ChatDatabase) -> None:
|
|
"""Test the generation time table keeps only the most recent entries."""
|
|
import sqlite3
|
|
|
|
from vibe_bot.database import IMAGE_GEN_TIME_LIMIT
|
|
|
|
total = IMAGE_GEN_TIME_LIMIT + 5
|
|
for i in range(total):
|
|
chat_db.record_image_generation_time(float(i))
|
|
|
|
conn = sqlite3.connect(chat_db.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT COUNT(*) FROM image_generation_times")
|
|
count = cursor.fetchone()[0]
|
|
cursor.execute("SELECT MIN(id) FROM image_generation_times")
|
|
min_id = cursor.fetchone()[0]
|
|
cursor.execute("SELECT MAX(id) FROM image_generation_times")
|
|
max_id = cursor.fetchone()[0]
|
|
conn.close()
|
|
|
|
assert count == IMAGE_GEN_TIME_LIMIT
|
|
assert min_id == 6
|
|
assert max_id == total
|
|
|
|
|
|
def test_image_generation_estimate_after_capping(
|
|
chat_db: ChatDatabase,
|
|
) -> None:
|
|
"""Test the estimate is computed from the retained recent rows."""
|
|
from vibe_bot.database import IMAGE_GEN_TIME_LIMIT
|
|
|
|
total = IMAGE_GEN_TIME_LIMIT + 5
|
|
for i in range(total):
|
|
chat_db.record_image_generation_time(float(i))
|
|
|
|
# Rows 95.0 through 104.0 are the ten most recent retained generations.
|
|
estimate = chat_db.get_image_generation_time_estimate()
|
|
assert estimate == pytest.approx(99.5)
|
|
|
|
|
|
def _broken_connection() -> MagicMock:
|
|
"""A mock connection whose first cursor.execute raises."""
|
|
fake_conn = MagicMock()
|
|
fake_conn.cursor.return_value.execute.side_effect = Exception("db error")
|
|
return fake_conn
|
|
|
|
|
|
def test_record_image_generation_time_failure(chat_db: ChatDatabase) -> None:
|
|
"""A database error while recording yields False, not an exception."""
|
|
with patch("vibe_bot.db.timing.connect", return_value=_broken_connection()):
|
|
assert chat_db.record_image_generation_time(1.0) is False
|
|
|
|
|
|
def test_image_generation_estimate_failure(chat_db: ChatDatabase) -> None:
|
|
"""A database error while reading yields None, not an exception."""
|
|
with patch("vibe_bot.db.timing.connect", return_value=_broken_connection()):
|
|
assert chat_db.get_image_generation_time_estimate() is None
|
|
|
|
|
|
def test_get_user_history(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""Test retrieving user message history."""
|
|
chat_db.add_message(
|
|
message_id="msg-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="User question",
|
|
)
|
|
chat_db.add_message(
|
|
message_id="msg-1_response",
|
|
user_id="bot",
|
|
username="vibe-bot",
|
|
content="Bot answer",
|
|
)
|
|
|
|
conversations = chat_db.get_user_history("u1")
|
|
assert len(conversations) == 1
|
|
assert conversations[0][0] == "User question"
|
|
assert conversations[0][1] == "Bot answer"
|
|
|
|
|
|
def test_get_user_history_no_response(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""Test user history when there is no bot response."""
|
|
chat_db.add_message(
|
|
message_id="msg-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="User question with no response",
|
|
)
|
|
|
|
conversations = chat_db.get_user_history("u1")
|
|
assert len(conversations) == 0
|
|
|
|
|
|
def test_get_user_history_excludes_bot(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""Test that bot messages are excluded from user history."""
|
|
chat_db.add_message(
|
|
message_id="msg-1",
|
|
user_id="bot",
|
|
username="vibe-bot",
|
|
content="Bot message",
|
|
)
|
|
|
|
conversations = chat_db.get_user_history("u1")
|
|
assert len(conversations) == 0
|
|
|
|
|
|
def test_get_bot_history(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""get_bot_history pairs user messages with responses for one bot."""
|
|
chat_db.add_message(
|
|
message_id="bh-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="bot question",
|
|
bot_name="alfred",
|
|
)
|
|
chat_db.add_message(
|
|
message_id="bh-1_response",
|
|
user_id="bot-1",
|
|
username="some-bot",
|
|
content="bot answer",
|
|
bot_name="alfred",
|
|
role="assistant",
|
|
embed=False,
|
|
)
|
|
chat_db.add_message(
|
|
message_id="bh-2",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="unanswered question",
|
|
bot_name="alfred",
|
|
)
|
|
|
|
history = chat_db.get_bot_history("alfred")
|
|
assert history == [("bot question", "bot answer")]
|
|
|
|
|
|
def test_get_conversation_context(
|
|
chat_db: ChatDatabase,
|
|
mock_embedding: MagicMock,
|
|
) -> None:
|
|
"""Test getting conversation context for RAG."""
|
|
chat_db.add_message(
|
|
message_id="msg-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="Previous question",
|
|
)
|
|
chat_db.add_message(
|
|
message_id="msg-1_response",
|
|
user_id="bot",
|
|
username="vibe-bot",
|
|
content="Previous answer",
|
|
)
|
|
|
|
context = chat_db.get_conversation_context("u1", "current message")
|
|
assert isinstance(context, list)
|
|
assert len(context) >= 2
|
|
|
|
|
|
def test_get_conversation_context_empty(chat_db: ChatDatabase) -> None:
|
|
"""Test getting context when there is no history."""
|
|
context = chat_db.get_conversation_context("u1", "new message")
|
|
assert context == []
|
|
|
|
|
|
def _query_vector() -> list[float]:
|
|
"""A 8-dim query vector along the first axis."""
|
|
return [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
|
|
|
|
|
|
def _seed_search_rows(chat_db: ChatDatabase) -> list[tuple[str, str, list[float]]]:
|
|
"""Seed user/response rows with known embeddings (plus exclusion traps)."""
|
|
import sqlite3
|
|
|
|
rows: list[tuple[str, str, list[float]]] = [
|
|
("s-1", "ask about the sky", [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
|
|
("s-2", "ask about the sea", [0.7, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
|
|
("s-3", "ask about the sun", [0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
|
|
("s-4", "ask about the sand", [0.2, 0.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
|
|
]
|
|
|
|
def stored_norm(vector: list[float]) -> float:
|
|
return float(np.linalg.norm(np.array(vector, dtype=np.float32)))
|
|
|
|
conn = sqlite3.connect(chat_db.db_path)
|
|
cursor = conn.cursor()
|
|
for i, (message_id, content, vector) in enumerate(rows):
|
|
timestamp = f"2024-01-0{i + 1} 00:00:00"
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages "
|
|
"(message_id, user_id, username, content, role, timestamp) "
|
|
"VALUES (?, ?, ?, ?, 'user', ?)",
|
|
(message_id, "u1", "alice", content, timestamp),
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages "
|
|
"(message_id, user_id, username, content, role, timestamp) "
|
|
"VALUES (?, ?, ?, ?, 'assistant', ?)",
|
|
(
|
|
f"{message_id}_response",
|
|
"bot-1",
|
|
"some-bot",
|
|
f"response {i + 1}",
|
|
f"2024-01-0{i + 1} 00:00:01",
|
|
),
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO message_embeddings (message_id, embedding, norm) "
|
|
"VALUES (?, ?, ?)",
|
|
(message_id, chat_db._vector_to_bytes(vector), stored_norm(vector)),
|
|
)
|
|
|
|
# A user row with the top possible similarity but no response row: the
|
|
# JOIN must exclude it instead of returning a NULL response.
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages "
|
|
"(message_id, user_id, username, content, role, timestamp) "
|
|
"VALUES ('s-5', 'u1', 'alice', 'orphan question', 'user', "
|
|
"'2024-01-05 00:00:00')",
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO message_embeddings (message_id, embedding, norm) "
|
|
"VALUES (?, ?, ?)",
|
|
(
|
|
"s-5",
|
|
chat_db._vector_to_bytes(_query_vector()),
|
|
stored_norm(_query_vector()),
|
|
),
|
|
)
|
|
|
|
# An assistant row carrying an embedding: the role filter must skip it.
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages "
|
|
"(message_id, user_id, username, content, role, timestamp) "
|
|
"VALUES ('s-6', 'bot-1', 'some-bot', 'assistant noise', 'assistant', "
|
|
"'2024-01-06 00:00:00')",
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO message_embeddings (message_id, embedding, norm) "
|
|
"VALUES (?, ?, ?)",
|
|
(
|
|
"s-6",
|
|
chat_db._vector_to_bytes(_query_vector()),
|
|
stored_norm(_query_vector()),
|
|
),
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return rows
|
|
|
|
|
|
def test_search_matches_reference_topk_and_ordering(
|
|
chat_db: ChatDatabase,
|
|
) -> None:
|
|
"""The JOINed search matches the per-row reference: same top-k and order."""
|
|
rows = _seed_search_rows(chat_db)
|
|
query = _query_vector()
|
|
|
|
expected: list[tuple[str, str, float]] = []
|
|
query_arr = np.array(query, dtype=np.float32)
|
|
for i, (_message_id, content, vector) in enumerate(rows):
|
|
stored = np.frombuffer(chat_db._vector_to_bytes(vector), dtype=np.float32)
|
|
similarity = float(
|
|
np.dot(query_arr, stored)
|
|
/ (np.linalg.norm(query_arr) * np.linalg.norm(stored))
|
|
)
|
|
expected.append((content, f"response {i + 1}", similarity))
|
|
expected.sort(key=lambda item: item[2], reverse=True)
|
|
expected_top = expected[:3]
|
|
|
|
with patch("vibe_bot.llm_client.embedding", return_value=query):
|
|
results = chat_db.search_similar_messages(
|
|
"query text", top_k=3, min_similarity=0.0
|
|
)
|
|
|
|
assert len(results) == 3
|
|
for (content, response, actual), (_ec, er, reference) in zip(
|
|
results, expected_top, strict=True
|
|
):
|
|
assert content == _ec
|
|
assert response == er
|
|
assert actual == pytest.approx(reference, abs=1e-5)
|
|
|
|
|
|
def test_search_with_stored_norms_matches_pre_norm_reference(
|
|
temp_db_path: str,
|
|
) -> None:
|
|
"""Stored-norm search is identical to per-vector renormalization.
|
|
|
|
Seeds a pre-norm database, migrates it, then asserts the stored-norm
|
|
search (top-k + ordering + similarities) matches a reference that
|
|
renormalizes every candidate vector inline — the pre-optimization
|
|
algorithm.
|
|
"""
|
|
from vibe_bot.database import ChatDatabase
|
|
|
|
rows = _seed_pre_norm_db(temp_db_path)
|
|
db = ChatDatabase(db_path=temp_db_path)
|
|
|
|
query = _query_vector()
|
|
with patch("vibe_bot.llm_client.embedding", return_value=query):
|
|
results = db.search_similar_messages("query text", top_k=3, min_similarity=0.0)
|
|
|
|
query_arr = np.array(query, dtype=np.float32)
|
|
query_norm = float(np.linalg.norm(query_arr))
|
|
expected: list[tuple[str, str, float]] = []
|
|
for i, (_message_id, content, vector) in enumerate(rows):
|
|
stored = np.array(vector, dtype=np.float32)
|
|
stored_norm = float(np.linalg.norm(stored))
|
|
similarity = (
|
|
0.0
|
|
if stored_norm == 0
|
|
else float(np.dot(query_arr, stored) / (query_norm * stored_norm))
|
|
)
|
|
expected.append((content, f"response {i + 1}", similarity))
|
|
expected.sort(key=lambda item: item[2], reverse=True)
|
|
|
|
assert len(results) == 3
|
|
for (content, response, actual), (ec, er, reference) in zip(
|
|
results, expected[:3], strict=True
|
|
):
|
|
assert content == ec
|
|
assert response == er
|
|
assert actual == pytest.approx(reference, abs=1e-6)
|
|
|
|
|
|
def test_search_null_norm_row_scores_zero(temp_db_path: str) -> None:
|
|
"""A row whose norm was never written scores 0 instead of crashing."""
|
|
import sqlite3
|
|
|
|
from vibe_bot.database import ChatDatabase
|
|
|
|
_rows = _seed_pre_norm_db(temp_db_path)
|
|
db = ChatDatabase(db_path=temp_db_path)
|
|
|
|
conn = sqlite3.connect(temp_db_path)
|
|
conn.execute("UPDATE message_embeddings SET norm = NULL WHERE message_id = 'n-1'")
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
query = _query_vector()
|
|
with patch("vibe_bot.llm_client.embedding", return_value=query):
|
|
results = db.search_similar_messages("query text", top_k=10, min_similarity=0.0)
|
|
|
|
scores = {content: similarity for content, _response, similarity in results}
|
|
assert scores["ask about the sky"] == 0.0
|
|
# All four scored rows plus the zero-vector row (0.0 passes min_similarity=0.0).
|
|
assert len(results) == 5
|
|
|
|
|
|
def test_search_mixed_embedding_dims_falls_back_to_per_row(temp_db_path: str) -> None:
|
|
"""Rows with different blob lengths (mid-life model change) don't crash."""
|
|
import sqlite3
|
|
|
|
from vibe_bot.database import ChatDatabase
|
|
|
|
wide = [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
|
|
narrow = [0.5, 0.5, 0.0, 0.0]
|
|
|
|
db = ChatDatabase(db_path=temp_db_path)
|
|
conn = sqlite3.connect(temp_db_path)
|
|
cursor = conn.cursor()
|
|
for message_id, content, vector in (
|
|
("m-1", "wide question", wide),
|
|
("m-2", "narrow question", narrow),
|
|
):
|
|
blob = np.array(vector, dtype=np.float32).tobytes()
|
|
stored_norm = float(np.linalg.norm(np.frombuffer(blob, dtype=np.float32)))
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages "
|
|
"(message_id, user_id, username, content, role, timestamp) "
|
|
"VALUES (?, 'u1', 'alice', ?, 'user', '2024-01-01 00:00:00')",
|
|
(message_id, content),
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO chat_messages "
|
|
"(message_id, user_id, username, content, role, timestamp) "
|
|
"VALUES (?, 'bot-1', 'some-bot', ?, 'assistant', '2024-01-01 00:00:01')",
|
|
(f"{message_id}_response", f"{message_id} answer"),
|
|
)
|
|
cursor.execute(
|
|
"INSERT INTO message_embeddings (message_id, embedding, norm) "
|
|
"VALUES (?, ?, ?)",
|
|
(message_id, blob, stored_norm),
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
query = _query_vector()
|
|
with patch("vibe_bot.llm_client.embedding", return_value=query):
|
|
results = db.search_similar_messages("query text", top_k=10, min_similarity=0.0)
|
|
|
|
scores = {content: similarity for content, _response, similarity in results}
|
|
# Both rows are returned: the 8-dim row scores exactly, the 4-dim row is
|
|
# zero-padded to the query dim and still scores sanely.
|
|
assert set(scores) == {"wide question", "narrow question"}
|
|
wide_norm = float(np.linalg.norm(np.array(wide, dtype=np.float32)))
|
|
narrow_norm = float(np.linalg.norm(np.array(narrow, dtype=np.float32)))
|
|
assert scores["wide question"] == pytest.approx(0.9 / wide_norm, abs=1e-5)
|
|
assert scores["narrow question"] == pytest.approx(0.5 / narrow_norm, abs=1e-5)
|
|
assert scores["wide question"] > scores["narrow question"]
|
|
|
|
|
|
def test_search_issues_single_select_per_call(chat_db: ChatDatabase) -> None:
|
|
"""search_similar_messages issues exactly one SELECT per call (no N+1)."""
|
|
import vibe_bot.db.search as db_search
|
|
from vibe_bot.db.connection import connect as realconnect
|
|
|
|
chat_db.add_message(
|
|
message_id="n-1",
|
|
user_id="u1",
|
|
username="alice",
|
|
content="one question",
|
|
)
|
|
chat_db.add_message(
|
|
message_id="n-1_response",
|
|
user_id="bot-1",
|
|
username="some-bot",
|
|
content="one answer",
|
|
role="assistant",
|
|
embed=False,
|
|
)
|
|
|
|
select_statements: list[str] = []
|
|
|
|
class _TracingCursor:
|
|
def __init__(self, cursor: sqlite3.Cursor) -> None:
|
|
self._cursor = cursor
|
|
|
|
def execute(self, sql: str, *args: Any) -> Any:
|
|
if sql.strip().upper().startswith("SELECT"):
|
|
select_statements.append(sql)
|
|
return self._cursor.execute(sql, *args)
|
|
|
|
def fetchall(self) -> Any:
|
|
return self._cursor.fetchall()
|
|
|
|
def fetchone(self) -> Any:
|
|
return self._cursor.fetchone()
|
|
|
|
class _TracingConnection:
|
|
def __init__(self, conn: sqlite3.Connection) -> None:
|
|
self._conn = conn
|
|
|
|
def cursor(self) -> _TracingCursor:
|
|
return _TracingCursor(self._conn.cursor())
|
|
|
|
def close(self) -> None:
|
|
self._conn.close()
|
|
|
|
def tracingconnect(db_path: str) -> _TracingConnection:
|
|
return _TracingConnection(realconnect(db_path))
|
|
|
|
with patch.object(db_search, "connect", side_effect=tracingconnect):
|
|
results = chat_db.search_similar_messages(
|
|
"one question", top_k=5, min_similarity=0.0
|
|
)
|
|
|
|
assert len(results) == 1
|
|
assert results[0][0] == "one question"
|
|
assert results[0][1] == "one answer"
|
|
assert len(select_statements) == 1
|
|
|
|
|
|
def test_search_empty_query_embedding_returns_empty(chat_db: ChatDatabase) -> None:
|
|
"""A failed query embedding yields no results."""
|
|
with patch("vibe_bot.llm_client.embedding", return_value=[]):
|
|
assert chat_db.search_similar_messages("anything") == []
|
|
|
|
|
|
def test_search_zero_query_vector_returns_empty(chat_db: ChatDatabase) -> None:
|
|
"""A zero-norm query vector yields no results (no division by zero)."""
|
|
with patch("vibe_bot.llm_client.embedding", return_value=[0.0] * 8):
|
|
assert chat_db.search_similar_messages("anything") == []
|
|
|
|
|
|
def test_custom_bot_create(custom_bot_manager: Any) -> None:
|
|
"""Test creating a custom bot returns "created" for a new name."""
|
|
result = custom_bot_manager.create_custom_bot(
|
|
bot_name="alfred",
|
|
system_prompt="You are a british butler",
|
|
created_by="user-123",
|
|
)
|
|
assert result == "created"
|
|
|
|
|
|
def test_custom_bot_create_duplicate(
|
|
custom_bot_manager: Any,
|
|
) -> None:
|
|
"""Test creating a duplicate custom bot replaces the old one."""
|
|
first = custom_bot_manager.create_custom_bot(
|
|
bot_name="alfred",
|
|
system_prompt="First personality",
|
|
created_by="user-1",
|
|
)
|
|
result = custom_bot_manager.create_custom_bot(
|
|
bot_name="alfred",
|
|
system_prompt="Second personality",
|
|
created_by="user-1",
|
|
)
|
|
assert first == "created"
|
|
assert result == "replaced"
|
|
|
|
bot = custom_bot_manager.get_custom_bot("alfred")
|
|
assert bot is not None
|
|
assert bot[1] == "Second personality"
|
|
|
|
|
|
def test_custom_bot_create_failure(custom_bot_manager: Any) -> None:
|
|
"""A database error while creating yields False, not an exception."""
|
|
with patch("vibe_bot.db.bots.connect", return_value=_broken_connection()):
|
|
result = custom_bot_manager.create_custom_bot(
|
|
bot_name="failbot",
|
|
system_prompt="a long enough personality",
|
|
created_by="user-1",
|
|
)
|
|
assert result is False
|
|
|
|
|
|
def test_custom_bot_create_case_insensitive(
|
|
custom_bot_manager: Any,
|
|
) -> None:
|
|
"""Test that bot names are case-insensitive."""
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="Alfred",
|
|
system_prompt="British butler",
|
|
created_by="user-1",
|
|
)
|
|
bot = custom_bot_manager.get_custom_bot("alfred")
|
|
assert bot is not None
|
|
|
|
|
|
def test_custom_bot_get_not_found(custom_bot_manager: Any) -> None:
|
|
"""Test getting a non-existent custom bot returns None."""
|
|
result = custom_bot_manager.get_custom_bot("nonexistent")
|
|
assert result is None
|
|
|
|
|
|
def test_custom_bot_get_returns_datetime(custom_bot_manager: Any) -> None:
|
|
"""created_at comes back as a real datetime, not a string."""
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="dtbot",
|
|
system_prompt="a long enough personality",
|
|
created_by="user-1",
|
|
)
|
|
result = custom_bot_manager.get_custom_bot("dtbot")
|
|
assert result is not None
|
|
assert isinstance(result[3], datetime)
|
|
|
|
|
|
def test_custom_bot_get_returns_correct_data(
|
|
custom_bot_manager: Any,
|
|
) -> None:
|
|
"""Test that get_custom_bot returns the correct bot data."""
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="testbot",
|
|
system_prompt="test prompt",
|
|
created_by="creator-1",
|
|
)
|
|
result = custom_bot_manager.get_custom_bot("testbot")
|
|
assert result is not None
|
|
assert result[0] == "testbot"
|
|
assert result[1] == "test prompt"
|
|
assert result[2] == "creator-1"
|
|
assert isinstance(result[3], datetime)
|
|
|
|
|
|
def test_custom_bot_list_empty(custom_bot_manager: Any) -> None:
|
|
"""Test listing custom bots when none exist."""
|
|
bots = custom_bot_manager.list_custom_bots()
|
|
assert bots == []
|
|
|
|
|
|
def test_custom_bot_list(custom_bot_manager: Any) -> None:
|
|
"""Test listing custom bots."""
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="bot-a",
|
|
system_prompt="prompt a",
|
|
created_by="user-1",
|
|
)
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="bot-b",
|
|
system_prompt="prompt b",
|
|
created_by="user-2",
|
|
)
|
|
|
|
bots = custom_bot_manager.list_custom_bots()
|
|
assert len(bots) == 2
|
|
|
|
|
|
def test_custom_bot_list_by_creator(custom_bot_manager: Any) -> None:
|
|
"""list_custom_bots filters by creator when user_id is given."""
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="bot-x",
|
|
system_prompt="prompt x",
|
|
created_by="user-1",
|
|
)
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="bot-y",
|
|
system_prompt="prompt y",
|
|
created_by="user-2",
|
|
)
|
|
|
|
bots = custom_bot_manager.list_custom_bots(user_id="user-1")
|
|
assert [bot[0] for bot in bots] == ["bot-x"]
|
|
|
|
|
|
def test_custom_bot_delete(custom_bot_manager: Any) -> None:
|
|
"""Test deleting a custom bot."""
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="deleteme",
|
|
system_prompt="will be deleted",
|
|
created_by="user-1",
|
|
)
|
|
result = custom_bot_manager.delete_custom_bot("deleteme")
|
|
assert result is True
|
|
|
|
bot = custom_bot_manager.get_custom_bot("deleteme")
|
|
assert bot is None
|
|
|
|
|
|
def test_custom_bot_delete_nonexistent(
|
|
custom_bot_manager: Any,
|
|
) -> None:
|
|
"""Test deleting a non-existent bot returns False."""
|
|
result = custom_bot_manager.delete_custom_bot("nonexistent")
|
|
assert result is False
|
|
|
|
|
|
def test_custom_bot_delete_failure(custom_bot_manager: Any) -> None:
|
|
"""A database error while deleting yields False, not an exception."""
|
|
with patch("vibe_bot.db.bots.connect", return_value=_broken_connection()):
|
|
assert custom_bot_manager.delete_custom_bot("whatever") is False
|
|
|
|
|
|
def test_custom_bot_list_excludes_inactive(
|
|
custom_bot_manager: Any,
|
|
) -> None:
|
|
"""Test that list_custom_bots excludes bots with is_active = 0."""
|
|
import sqlite3
|
|
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="active-bot",
|
|
system_prompt="stays active",
|
|
created_by="user-1",
|
|
)
|
|
custom_bot_manager.create_custom_bot(
|
|
bot_name="deactivated-bot",
|
|
system_prompt="should not appear",
|
|
created_by="user-1",
|
|
)
|
|
conn = sqlite3.connect(custom_bot_manager.db_path)
|
|
conn.execute(
|
|
"UPDATE custom_bots SET is_active = 0 WHERE bot_name = 'deactivated-bot'"
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
bots = custom_bot_manager.list_custom_bots()
|
|
assert len(bots) == 1
|
|
assert bots[0][0] == "active-bot"
|
|
|
|
|
|
def test_custom_bot_delete_with_error(
|
|
custom_bot_manager: Any,
|
|
) -> None:
|
|
"""Test that delete_custom_bot returns False on error."""
|
|
with patch.object(
|
|
custom_bot_manager,
|
|
"_initialize_custom_bots_table",
|
|
side_effect=Exception("db error"),
|
|
):
|
|
pass
|
|
result = custom_bot_manager.delete_custom_bot("nonexistent")
|
|
assert result is False
|
|
|
|
|
|
def test_database_get_database_singleton(temp_db_path: str) -> None:
|
|
"""Test that get_database returns the same instance."""
|
|
import vibe_bot.database as db_module
|
|
from vibe_bot.database import ChatDatabase, get_database
|
|
|
|
db_module._chat_db = None
|
|
|
|
db1 = get_database()
|
|
assert isinstance(db1, ChatDatabase)
|
|
|
|
db2 = get_database()
|
|
assert db1 is db2
|
|
|
|
|
|
def test_database_init_creates_tables(temp_db_path: str) -> None:
|
|
"""Test that database initialization creates the expected tables."""
|
|
from vibe_bot.database import ChatDatabase, CustomBotManager
|
|
|
|
ChatDatabase(db_path=temp_db_path)
|
|
CustomBotManager(db_path=temp_db_path)
|
|
|
|
import sqlite3
|
|
|
|
conn = sqlite3.connect(temp_db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
|
tables = {row[0] for row in cursor.fetchall()}
|
|
conn.close()
|
|
|
|
assert "chat_messages" in tables
|
|
assert "message_embeddings" in tables
|
|
assert "custom_bots" in tables
|
|
assert "image_generation_times" in tables
|