From c1ba742c0d5c0c354fc82cce3b29de3903f5a8e6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Mon, 30 Oct 2023 19:17:43 +0100
Subject: [PATCH] test: added test for is_supported_codeblock

---
 tests/test_util.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)
 create mode 100644 tests/test_util.py

diff --git a/tests/test_util.py b/tests/test_util.py
new file mode 100644
index 00000000..73f7c5b1
--- /dev/null
+++ b/tests/test_util.py
@@ -0,0 +1,25 @@
+from gptme.tools import is_supported_codeblock
+
+
+def test_is_supported_codeblock():
+    block_plain = """```
+some plaintext
+```
+"""
+    # clean unsupported block
+    assert not is_supported_codeblock(block_plain)
+
+    block_python = """```python
+print("hello world")
+```
+"""
+    # clean supported block
+    assert is_supported_codeblock(block_python)
+
+    # has preamble
+    s = f"""bla bla\n{block_python}"""
+    assert is_supported_codeblock(s)
+
+    # last block is plain/unsupported
+    s = f"""bla bla\n{block_python}\nbla bla{block_plain}"""
+    assert not is_supported_codeblock(s)