From ea52a3ee1c09bcae89eb2827468f4205d2243e54 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Mon, 18 Nov 2024 22:20:59 -0800
Subject: [PATCH] minor enhancement for test fixtures

---
 llama_stack/providers/tests/agents/fixtures.py | 2 +-
 llama_stack/providers/tests/safety/fixtures.py | 3 +++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/llama_stack/providers/tests/agents/fixtures.py b/llama_stack/providers/tests/agents/fixtures.py
index 1f89b909a..93a011c95 100644
--- a/llama_stack/providers/tests/agents/fixtures.py
+++ b/llama_stack/providers/tests/agents/fixtures.py
@@ -83,6 +83,6 @@ async def agents_stack(request, inference_model, safety_shield):
             )
             for model in inference_models
         ],
-        shields=[safety_shield],
+        shields=[safety_shield] if safety_shield else [],
     )
     return test_stack
diff --git a/llama_stack/providers/tests/safety/fixtures.py b/llama_stack/providers/tests/safety/fixtures.py
index a706316dd..32883bfab 100644
--- a/llama_stack/providers/tests/safety/fixtures.py
+++ b/llama_stack/providers/tests/safety/fixtures.py
@@ -47,6 +47,9 @@ def safety_shield(request):
     else:
         params = {}
 
+    if not shield_id:
+        return None
+
     return ShieldInput(
         shield_id=shield_id,
         params=params,