From 75480b01b8650d0fba3cb2aedaae75743b3dff67 Mon Sep 17 00:00:00 2001
From: Mustafa Elbehery <melbeher@redhat.com>
Date: Fri, 18 Jul 2025 11:54:02 +0200
Subject: [PATCH] chore(test): migrate unit tests from unittest to pytest for
 system prompt (#2789)

This PR replaces unittest with pytest.

Part of https://github.com/meta-llama/llama-stack/issues/2680

cc @leseb

Signed-off-by: Mustafa Elbehery <melbeher@redhat.com>
---
 tests/unit/models/test_system_prompts.py | 113 ++++++++++++-----------
 1 file changed, 57 insertions(+), 56 deletions(-)

diff --git a/tests/unit/models/test_system_prompts.py b/tests/unit/models/test_system_prompts.py
index 1f4ccc7e3..f5580f4c5 100644
--- a/tests/unit/models/test_system_prompts.py
+++ b/tests/unit/models/test_system_prompts.py
@@ -12,7 +12,6 @@
 # the top-level of this source tree.
 
 import textwrap
-import unittest
 from datetime import datetime
 
 from llama_stack.models.llama.llama3.prompt_templates import (
@@ -24,59 +23,61 @@ from llama_stack.models.llama.llama3.prompt_templates import (
 )
 
 
-class PromptTemplateTests(unittest.TestCase):
-    def check_generator_output(self, generator):
-        for example in generator.data_examples():
-            pt = generator.gen(example)
-            text = pt.render()
-            # print(text)  # debugging
-            if not example:
-                continue
-            for tool in example:
-                assert tool.tool_name in text
-
-    def test_system_default(self):
-        generator = SystemDefaultGenerator()
-        today = datetime.now().strftime("%d %B %Y")
-        expected_text = f"Cutting Knowledge Date: December 2023\nToday Date: {today}"
-        assert expected_text.strip("\n") == generator.gen(generator.data_examples()[0]).render()
-
-    def test_system_builtin_only(self):
-        generator = BuiltinToolGenerator()
-        expected_text = textwrap.dedent(
-            """
-            Environment: ipython
-            Tools: brave_search, wolfram_alpha
-            """
-        )
-        assert expected_text.strip("\n") == generator.gen(generator.data_examples()[0]).render()
-
-    def test_system_custom_only(self):
-        self.maxDiff = None
-        generator = JsonCustomToolGenerator()
-        self.check_generator_output(generator)
-
-    def test_system_custom_function_tag(self):
-        self.maxDiff = None
-        generator = FunctionTagCustomToolGenerator()
-        self.check_generator_output(generator)
-
-    def test_llama_3_2_system_zero_shot(self):
-        generator = PythonListCustomToolGenerator()
-        self.check_generator_output(generator)
-
-    def test_llama_3_2_provided_system_prompt(self):
-        generator = PythonListCustomToolGenerator()
-        user_system_prompt = textwrap.dedent(
-            """
-            Overriding message.
-
-            {{ function_description }}
-            """
-        )
-        example = generator.data_examples()[0]
-
-        pt = generator.gen(example, user_system_prompt)
+def check_generator_output(generator):
+    for example in generator.data_examples():
+        pt = generator.gen(example)
         text = pt.render()
-        assert "Overriding message." in text
-        assert '"name": "get_weather"' in text
+        if not example:
+            continue
+        for tool in example:
+            assert tool.tool_name in text
+
+
+def test_system_default():
+    generator = SystemDefaultGenerator()
+    today = datetime.now().strftime("%d %B %Y")
+    expected_text = f"Cutting Knowledge Date: December 2023\nToday Date: {today}"
+    assert expected_text.strip("\n") == generator.gen(generator.data_examples()[0]).render()
+
+
+def test_system_builtin_only():
+    generator = BuiltinToolGenerator()
+    expected_text = textwrap.dedent(
+        """
+        Environment: ipython
+        Tools: brave_search, wolfram_alpha
+        """
+    )
+    assert expected_text.strip("\n") == generator.gen(generator.data_examples()[0]).render()
+
+
+def test_system_custom_only():
+    generator = JsonCustomToolGenerator()
+    check_generator_output(generator)
+
+
+def test_system_custom_function_tag():
+    generator = FunctionTagCustomToolGenerator()
+    check_generator_output(generator)
+
+
+def test_llama_3_2_system_zero_shot():
+    generator = PythonListCustomToolGenerator()
+    check_generator_output(generator)
+
+
+def test_llama_3_2_provided_system_prompt():
+    generator = PythonListCustomToolGenerator()
+    user_system_prompt = textwrap.dedent(
+        """
+        Overriding message.
+
+        {{ function_description }}
+        """
+    )
+    example = generator.data_examples()[0]
+
+    pt = generator.gen(example, user_system_prompt)
+    text = pt.render()
+    assert "Overriding message." in text
+    assert '"name": "get_weather"' in text