From f0ba367877630989e7c8ff70d931e6b57df21922 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Fri, 31 Jan 2025 15:29:55 -0800
Subject: [PATCH] Update client-sdk test config option handling

---
 tests/client-sdk/conftest.py        | 6 ++++++
 tests/client-sdk/safety/conftest.py | 9 ---------
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/tests/client-sdk/conftest.py b/tests/client-sdk/conftest.py
index 2a366e985..8c44242fe 100644
--- a/tests/client-sdk/conftest.py
+++ b/tests/client-sdk/conftest.py
@@ -52,6 +52,12 @@ def pytest_addoption(parser):
         default=VISION_MODEL,
         help="Specify the vision inference model to use for testing",
     )
+    parser.addoption(
+        "--safety-shield",
+        action="store",
+        default="meta-llama/Llama-Guard-3-1B",
+        help="Specify the safety shield model to use for testing",
+    )
 
 
 @pytest.fixture(scope="session")
diff --git a/tests/client-sdk/safety/conftest.py b/tests/client-sdk/safety/conftest.py
index c4570801c..953b76cbf 100644
--- a/tests/client-sdk/safety/conftest.py
+++ b/tests/client-sdk/safety/conftest.py
@@ -5,15 +5,6 @@
 # the root directory of this source tree.
 
 
-def pytest_addoption(parser):
-    parser.addoption(
-        "--safety-shield",
-        action="store",
-        default="meta-llama/Llama-Guard-3-1B",
-        help="Specify the safety shield model to use for testing",
-    )
-
-
 def pytest_generate_tests(metafunc):
     if "llama_guard_text_shield_id" in metafunc.fixturenames:
         metafunc.parametrize(