From c721262ea3efa462b5ee83c3689be073d916fc5d Mon Sep 17 00:00:00 2001
From: Robert Szefler <robert@robusta.dev>
Date: Sat, 15 Jun 2024 15:30:55 +0200
Subject: [PATCH] trim_markdown: handle the case when there are no code blocks
 in the input; more unit tests

---
 src/robusta/utils/trim_markdown.py | 22 ++++++++++++++++++----
 tests/test_trim_markdown.py        | 26 +++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/src/robusta/utils/trim_markdown.py b/src/robusta/utils/trim_markdown.py
index 7bc42ea5d..10418b454 100644
--- a/src/robusta/utils/trim_markdown.py
+++ b/src/robusta/utils/trim_markdown.py
@@ -23,21 +23,35 @@ def trim_markdown(text: str, max_length: int, suffix: str = "...") -> str:
         if trim_idx >= close_end:
             # Trimming point after this block quote
             return text[:trim_idx] + suffix
-        if trim_idx < open_start:
+        elif trim_idx < open_start:
             # Trimming point before this block quote - continue to the preceding block
             continue
-        if trim_idx >= open_start and trim_idx < open_start + 3:
+        elif trim_idx >= open_start and trim_idx < open_start + 3:
             # Trimming point inside the opening block quote tag
             return text[:trim_idx].rstrip("`") + suffix
-        if trim_idx >= close_start and trim_idx < close_end:
+        elif trim_idx >= close_start and trim_idx < close_end:
             # Trimming point inside the closing block quote tag
             if trim_idx - open_end >= 3:  # Enough space to insert the closing tag
                 return text[:trim_idx - 3] + "```" + suffix
             else:  # Not enough space, strip the whole block
                 return text[:open_start] + suffix
-        if trim_idx >= open_end and trim_idx < close_start:
+        elif trim_idx >= open_end and trim_idx < close_start:
             # Trimming point inside the block quote
             if trim_idx - open_end >= 3:  # Enough space to insert the closing tag
                 return text[:trim_idx - 3] + "```" + suffix
             else:  # Not enough space, strip the whole block
                 return text[:open_start] + suffix
+        else:
+            # This should never happen
+            raise Exception(
+                f'Internal error in trim_markdown, text="{text[:12]}"(...), {max_length=}, suffix="{suffix}", '
+                f'matched code block {open_start}..{close_end}'
+            )
+
+    # Cases when there were no code blocks in the input
+    if len(text) <= trim_idx:
+        return text
+    elif len(text) < max_length:
+        return (text[:trim_idx] + suffix)[:max_length]
+    else:
+        return text[:trim_idx] + suffix
diff --git a/tests/test_trim_markdown.py b/tests/test_trim_markdown.py
index c725164cb..c708e9c32 100644
--- a/tests/test_trim_markdown.py
+++ b/tests/test_trim_markdown.py
@@ -42,6 +42,30 @@
     ])
 def test_trim_markdown(max_length: int, expected_output: str):
     text = "```oh``` hello ```world``` and then ```something```"
-    trimmed = trim_markdown(text, max_length, '##')
+    trimmed = trim_markdown(text, max_length, "##")
+    assert trimmed == expected_output
+    assert len(trimmed) <= max_length
+
+
+@pytest.mark.parametrize(
+    "max_length,expected_output", [
+        (0, ""),
+        (1, "$"),
+        (2, "$$"),
+        (3, "$$$"),
+        (4, "N$$$"),
+        (5, "No$$$"),
+        (10, "No code$$$"),
+        (38, "No code blocks whatsoever in this t$$$"),
+        (39, "No code blocks whatsoever in this te$$$"),
+        (40, "No code blocks whatsoever in this tex$$$"),
+        (41, "No code blocks whatsoever in this text"),
+        (42, "No code blocks whatsoever in this text"),
+        (111, "No code blocks whatsoever in this text"),
+    ]
+)
+def test_trim_markdown_no_code_blocks(max_length: int, expected_output: str):
+    text = "No code blocks whatsoever in this text"
+    trimmed = trim_markdown(text, max_length, "$$$")
     assert trimmed == expected_output
     assert len(trimmed) <= max_length