1 files changed, 79 insertions, 0 deletions
diff --git a/test_llamachat.py b/test_llamachat.py
index bb5f097..7a84608 100755
--- a/test_llamachat.py
+++ b/test_llamachat.py
@@ -441,6 +441,83 @@ def test_unbalanced_backticks():
     print("ok  unbalanced backticks")
 
 
+def test_code_block_wrapping():
+    """A long code line wraps instead of widening the whole transcript."""
+    import os
+
+    os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
+    from PySide6.QtGui import QTextDocument
+    from PySide6.QtWidgets import QApplication
+
+    from llamachat.ui import _markdown_to_fragment
+
+    app = QApplication.instance() or QApplication([])
+    assert app is not None
+
+    # Qt's markdown stylesheet wraps p and li but not pre, and it is dropped
+    # when only the body fragment is kept, so pre has to say so itself.
+    fragment = _markdown_to_fragment("```\n" + "A" * 200 + "\n```\n")
+    assert "white-space:pre-wrap" in fragment, fragment
+
+    # The real check is the layout: no block may exceed the viewport.
+    doc = QTextDocument()
+    doc.setHtml(fragment)
+    doc.setTextWidth(300)
+    layout = doc.documentLayout()
+    widest = 0.0
+    block = doc.begin()
+    while block.isValid():
+        widest = max(widest, layout.blockBoundingRect(block).width())
+        block = block.next()
+    assert widest <= 300, widest
+    print("ok  code block wrapping")
+
+
+def test_code_block_copy_links():
+    """Each fenced block gets a copy link addressing it by index."""
+    import os
+    import re as _re
+
+    os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
+    from PySide6.QtWidgets import QApplication
+
+    from llamachat.ui import (
+        COPY_SCHEME, _close_fences, _code_blocks, _markdown_to_fragment,
+    )
+
+    app = QApplication.instance() or QApplication([])
+    assert app is not None
+
+    # Blocks are read from the source: Qt emits one <pre> per line, so the
+    # text is not recoverable from the rendered fragment.
+    assert _code_blocks("```py\nx=1\ny=2\n```\n") == ["x=1\ny=2"]
+    assert _code_blocks("~~~\na\n~~~\n") == ["a"]
+    assert _code_blocks("no code here") == []
+    # A longer fence can hold shorter ones without ending the block.
+    assert _code_blocks("````\n```\ninner\n```\n````\n") == ["```\ninner\n```"]
+
+    # A link per block, numbered bubble.block so several replies coexist.
+    markdown = "one\n\n```\nAAA\n```\n\ntwo\n\n```\nBBB\n```\n"
+    fragment = _markdown_to_fragment(markdown, 3)
+    assert _re.findall(r"x-llamachat-copy:([0-9.]+)", fragment) == ["3.0", "3.1"]
+
+    # The indices the links carry must select the blocks the parser found.
+    blocks = _code_blocks(_close_fences(markdown))
+    assert blocks == ["AAA", "BBB"], blocks
+
+    # A block still streaming is closed first, so it is copyable mid-reply.
+    assert _code_blocks(_close_fences("```\nhalf")) == ["half"]
+
+    # Without an index there are no links, so a bubble that cannot be
+    # addressed does not emit a link that would not resolve.
+    assert COPY_SCHEME not in _markdown_to_fragment(markdown)
+
+    # A reply cannot forge one: the scheme is defused like the others.
+    forged = _markdown_to_fragment(f"[copy]({COPY_SCHEME}0.0)", 0)
+    assert f'href="{COPY_SCHEME}' not in forged, forged
+    print("ok  code block copy links")
+
+
 def test_system_qt_theme_guard():
     """The plugin path is only borrowed when the Qt versions agree."""
     import os
@@ -3260,4 +3337,6 @@ if __name__ == "__main__":
     test_title_request()
     test_needs_title()
     test_title_prompt()
+    test_code_block_wrapping()
+    test_code_block_copy_links()
     print("\nall checks passed")