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

vibe_bot/tests/test_main.py

@@ -1022,6 +1022,7 @@ def test_doodlebob_selects_portrait(
     mock_ctx: MagicMock,
     mock_llama_wrapper: MagicMock,
     mock_base64: MagicMock,
+    mock_database: MagicMock,
 ) -> None:
     """Test doodlebob picks portrait and passes the portrait size."""
     import asyncio
@@ -1033,6 +1034,7 @@ def test_doodlebob_selects_portrait(
         "a tall portrait of a lighthouse",  # prompt rewrite
     ]
     mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+    mock_database.get_image_generation_time_estimate.return_value = None
 
     with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
         asyncio.run(main_module.doodlebob(mock_ctx, message="a tall lighthouse"))
@@ -1055,6 +1057,7 @@ def test_doodlebob_selects_landscape(
     mock_ctx: MagicMock,
     mock_llama_wrapper: MagicMock,
     mock_base64: MagicMock,
+    mock_database: MagicMock,
 ) -> None:
     """Test doodlebob picks landscape and passes the landscape size."""
     import asyncio
@@ -1066,6 +1069,7 @@ def test_doodlebob_selects_landscape(
         "a wide panoramic coastline",
     ]
     mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+    mock_database.get_image_generation_time_estimate.return_value = None
 
     with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
         asyncio.run(main_module.doodlebob(mock_ctx, message="wide coastline"))
@@ -1077,6 +1081,7 @@ def test_doodlebob_malformed_layout_defaults_square(
     mock_ctx: MagicMock,
     mock_llama_wrapper: MagicMock,
     mock_base64: MagicMock,
+    mock_database: MagicMock,
 ) -> None:
     """Test a malformed layout response falls back to the square size."""
     import asyncio
@@ -1088,6 +1093,7 @@ def test_doodlebob_malformed_layout_defaults_square(
         "a balanced composition",
     ]
     mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+    mock_database.get_image_generation_time_estimate.return_value = None
 
     with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
         asyncio.run(main_module.doodlebob(mock_ctx, message="a logo"))
@@ -1101,6 +1107,7 @@ def test_doodlebob_empty_layout_defaults_square(
     mock_ctx: MagicMock,
     mock_llama_wrapper: MagicMock,
     mock_base64: MagicMock,
+    mock_database: MagicMock,
 ) -> None:
     """Test an empty layout response (LLM failure) falls back to square."""
     import asyncio
@@ -1112,8 +1119,96 @@ def test_doodlebob_empty_layout_defaults_square(
         "a balanced composition",
     ]
     mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+    mock_database.get_image_generation_time_estimate.return_value = None
 
     with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
         asyncio.run(main_module.doodlebob(mock_ctx, message="a logo"))
 
     assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1024"
+
+
+def test_doodlebob_reports_estimate_and_elapsed(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+    mock_database: MagicMock,
+) -> None:
+    """Doodlebob posts an ETA from history and the final elapsed time."""
+    import asyncio
+    import re
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "square",
+        "a test scene",
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+    mock_database.get_image_generation_time_estimate.return_value = 12.34
+
+    with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
+        asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
+
+    mock_database.record_image_generation_time.assert_called_once()
+    recorded = mock_database.record_image_generation_time.call_args[0][0]
+    assert isinstance(recorded, float)
+    assert recorded >= 0.0
+
+    sent = _sent_texts(mock_ctx)
+    assert any("drone strike" in m for m in sent)
+    assert any("~12 seconds" in m for m in sent)
+    assert any(re.search(r"generated in \d+\.\d+ seconds", m) for m in sent)
+
+
+def test_doodlebob_no_estimate_without_history(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+    mock_database: MagicMock,
+) -> None:
+    """No ETA is posted when there is no generation history yet."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "square",
+        "a test scene",
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+    mock_database.get_image_generation_time_estimate.return_value = None
+
+    with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
+        asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
+
+    sent = _sent_texts(mock_ctx)
+    assert any("drone strike" in m for m in sent)
+    assert not any("ETA" in m for m in sent)
+    mock_database.record_image_generation_time.assert_called_once()
+
+
+def test_doodlebob_failed_generation_not_recorded(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+    mock_database: MagicMock,
+) -> None:
+    """A failed generation records no time and reports no elapsed seconds."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "square",
+        "a test scene",
+    ]
+    mock_llama_wrapper.image_generation.return_value = ""
+    mock_database.get_image_generation_time_estimate.return_value = 8.0
+
+    with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
+        asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
+
+    mock_database.record_image_generation_time.assert_not_called()
+    sent = _sent_texts(mock_ctx)
+    assert any("Failed to generate image" in m for m in sent)
+    assert not any("generated in" in m for m in sent)