From bdb4396a9f85b16b9952100f19d419aed45e9d74 Mon Sep 17 00:00:00 2001
From: Eric Ma <ericmajinglong@gmail.com>
Date: Thu, 5 Sep 2024 20:32:08 -0400
Subject: [PATCH] =?UTF-8?q?refactor(test=5Fdocs)=F0=9F=94=A7:=20Refactor?=
 =?UTF-8?q?=20test=5Fdocs=20to=20use=20external=20source=20file=20and=20im?=
 =?UTF-8?q?prove=20function=20documentation?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Import pathlib and use Path for reading external prime number source code.
- Update docstrings in is_prime and next_prime with parameter and return details.
- Modify is_prime base case to check for n < 2 and add a case for n == 3.
---
 tests/cli/test_docs.py | 43 +++++++++++++++++-------------------------
 1 file changed, 17 insertions(+), 26 deletions(-)

diff --git a/tests/cli/test_docs.py b/tests/cli/test_docs.py
index f0318d3c..e7c9c34a 100644
--- a/tests/cli/test_docs.py
+++ b/tests/cli/test_docs.py
@@ -3,6 +3,8 @@
 This also serves as an evaluation suite for local LLMs that are tested.
 """
 
+from pathlib import Path
+
 import pytest
 from llamabot import StructuredBot
 from llamabot.cli.docs import (
@@ -85,31 +87,16 @@
 By combining these two functions, you can efficiently determine the primality of a number and find the next prime number with minimal computation.
 """
 
-original_source_code = '''
-def is_prime(n):
-    """Check if a number is prime."""
-    if n <= 1:
-        return False
-    if n == 2:
-        return True
-    if n % 2 == 0:
-        return False
-    for i in range(3, int(n**0.5) + 1, 2):
-        if n % i == 0:
-            return False
-    return True
-
-def next_prime(current_number):
-    """Find the next prime number after the current number."""
-    next_number = current_number + 1
-    while not is_prime(next_number):
-        next_number += 1
-    return next_number
-'''
+with open(Path(__file__).parent / "assets" / "next_prime" / "source.py", "r+") as f:
+    original_source_code = f.read()
 
 new_source_code = '''
 def is_prime(n):
-    """Check if a number is prime."""
+    """Check if a number is prime.
+
+    :param n: The number to check.
+    :return: True if the number is prime, False otherwise.
+    """
     if n < 2:  # Changed from n <= 1
         return False
     if n == 2 or n == 3:  # Added an additional base case for 3
@@ -124,7 +111,11 @@ def is_prime(n):
     return True
 
 def next_prime(current_number):
-    """Find the next prime number after the current number."""
+    """Find the next prime number after the current number.
+
+    :param current_number: The current number.
+    :return: The next prime number.
+    """
     next_number = current_number + 1
     while not is_prime(next_number):
         next_number += 1
@@ -199,10 +190,10 @@ def next_prime(current_number):
 @pytest.mark.parametrize(
     "model_name",
     [
-        # "ollama_chat/gemma2:2b",  # passes all tests
-        "gpt-4-turbo",  # passes all tests, but costs $$ to run
+        # "ollama_chat/gemma2:2b",  # does not pass all tests.
+        "gpt-4-turbo",  # passes all tests, but costs $$ to run.
         # ollama_chat/phi3", # garbage model, doesn't pass any tests.
-        "gpt-4o-mini",  # passes all tests, but costs $$ to run
+        "gpt-4o-mini",  # passes all tests, but costs $$ to run.
     ],
 )
 def test_out_of_date_when_source_changes(