From 05d1ead02f8ee2c3ff34be9fb89d9a5e6bf91e7a Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Tue, 19 Nov 2024 13:25:36 -0800
Subject: [PATCH] Update condition in tests to handle llama-3.1 vs llama3.1 (HF
 names)

---
 .../providers/tests/inference/test_text_inference.py      | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/llama_stack/providers/tests/inference/test_text_inference.py b/llama_stack/providers/tests/inference/test_text_inference.py
index 7b7aca5bd..6e263432a 100644
--- a/llama_stack/providers/tests/inference/test_text_inference.py
+++ b/llama_stack/providers/tests/inference/test_text_inference.py
@@ -25,7 +25,11 @@ from .utils import group_chunks
 
 
 def get_expected_stop_reason(model: str):
-    return StopReason.end_of_message if "Llama3.1" in model else StopReason.end_of_turn
+    return (
+        StopReason.end_of_message
+        if ("Llama3.1" in model or "Llama-3.1" in model)
+        else StopReason.end_of_turn
+    )
 
 
 @pytest.fixture
@@ -34,7 +38,7 @@ def common_params(inference_model):
         "tool_choice": ToolChoice.auto,
         "tool_prompt_format": (
             ToolPromptFormat.json
-            if "Llama3.1" in inference_model
+            if ("Llama3.1" in inference_model or "Llama-3.1" in inference_model)
             else ToolPromptFormat.python_list
         ),
     }