allow selection between portrait, landscape, square canvas types

vibe_bot/tests/test_main.py

@@ -961,3 +961,159 @@ def test_debug_tools(mock_ctx: MagicMock) -> None:
     assert "LLM Tools" in call_args
     assert "get_channel_members" in call_args
     assert "members" in call_args.lower()
+
+
+# ---------------------------------------------------------------------------
+# Image layout selection (doodlebob)
+# ---------------------------------------------------------------------------
+
+LAYOUT_TEST_SIZES = {
+    "portrait": "1024x1536",
+    "landscape": "1536x1024",
+    "square": "1024x1024",
+}
+
+
+def _sent_texts(mock_ctx: MagicMock) -> list[str]:
+    """Collect positional (text) arguments sent via ctx.send."""
+    return [c.args[0] for c in mock_ctx.send.call_args_list if c.args]
+
+
+@pytest.mark.parametrize(
+    ("response", "expected"),
+    [
+        ("portrait", "portrait"),
+        ("landscape", "landscape"),
+        ("square", "square"),
+        ("  Portrait  ", "portrait"),
+        ("LANDSCAPE", "landscape"),
+        ("square.", "square"),
+        ("I would use portrait.", "portrait"),
+        ("This scene is best as landscape.", "landscape"),
+        ("A square composition works here.", "square"),
+    ],
+)
+def test_parse_image_layout_valid(response: str, expected: str) -> None:
+    """Test that valid LLM layout responses parse to the right layout."""
+    import vibe_bot.main as main_module
+
+    assert main_module.parse_image_layout(response) == expected
+
+
+@pytest.mark.parametrize(
+    "response",
+    [
+        "",
+        "   ",
+        "banana",
+        "1024x1024",
+        "tall and wide",
+        "squareness",  # word boundary should prevent a match on "square"
+    ],
+)
+def test_parse_image_layout_defaults_to_square(response: str) -> None:
+    """Test that empty or malformed responses fall back to square."""
+    import vibe_bot.main as main_module
+
+    assert main_module.parse_image_layout(response) == "square"
+
+
+def test_doodlebob_selects_portrait(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+) -> None:
+    """Test doodlebob picks portrait and passes the portrait size."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "portrait",  # layout selection
+        "a tall portrait of a lighthouse",  # prompt rewrite
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+
+    with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
+        asyncio.run(main_module.doodlebob(mock_ctx, message="a tall lighthouse"))
+
+    # Layout selection is the first instruct call, using the layout prompt.
+    layout_call = mock_llama_wrapper.chat_completion_instruct.call_args_list[0]
+    assert layout_call.kwargs["system_prompt"] == main_module.IMAGE_LAYOUT_SYSTEM_PROMPT
+    assert layout_call.kwargs["user_prompt"] == "a tall lighthouse"
+
+    mock_llama_wrapper.image_generation.assert_called_once()
+    assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1536"
+
+    sent = _sent_texts(mock_ctx)
+    assert any("shopping for a canvas" in m for m in sent)
+    assert any("selected portrait" in m for m in sent)
+    assert any("drone strike" in m for m in sent)
+
+
+def test_doodlebob_selects_landscape(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+) -> None:
+    """Test doodlebob picks landscape and passes the landscape size."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "landscape",
+        "a wide panoramic coastline",
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+
+    with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
+        asyncio.run(main_module.doodlebob(mock_ctx, message="wide coastline"))
+
+    assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1536x1024"
+
+
+def test_doodlebob_malformed_layout_defaults_square(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+) -> None:
+    """Test a malformed layout response falls back to the square size."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "banana",  # malformed layout
+        "a balanced composition",
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+
+    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"
+    sent = _sent_texts(mock_ctx)
+    assert any("selected square" in m for m in sent)
+
+
+def test_doodlebob_empty_layout_defaults_square(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+) -> None:
+    """Test an empty layout response (LLM failure) falls back to square."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "",  # empty layout response
+        "a balanced composition",
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+
+    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"