add time estimation for image gen

vibe_bot/tests/test_database.py

@@ -194,6 +194,76 @@ def test_clear_all_messages(
     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 test_get_user_history(
     chat_db: ChatDatabase,
     mock_embedding: MagicMock,
@@ -484,3 +554,4 @@ def test_database_init_creates_tables(temp_db_path: str) -> None:
     assert "chat_messages" in tables
     assert "message_embeddings" in tables
     assert "custom_bots" in tables
+    assert "image_generation_times" in tables