7777from tests .helpers import new_worker
7878from tests .helpers .nexus import create_nexus_endpoint , make_nexus_endpoint_name
7979
80- response_index : int = 0
81-
8280
8381class StaticTestModel (TestModel ):
8482 __test__ = False
8583 responses : list [ModelResponse ] = []
8684
87- def response (self ):
88- global response_index
89- response = self .responses [response_index ]
90- response_index += 1
91- return response
92-
9385 def __init__ (
9486 self ,
9587 ) -> None :
96- global response_index
97- response_index = 0
98- super ().__init__ (self .response )
88+ self ._responses = iter (self .responses )
89+ super ().__init__ (lambda : next (self ._responses ))
9990
10091
10192class TestHelloModel (StaticTestModel ):
@@ -1348,9 +1339,6 @@ async def test_customer_service_workflow(client: Client, use_local_model: bool):
13481339 )
13491340
13501341
1351- guardrail_response_index : int = 0
1352-
1353-
13541342class InputGuardrailModel (OpenAIResponsesModel ):
13551343 __test__ = False
13561344 responses : list [ModelResponse ] = [
@@ -1439,11 +1427,9 @@ def __init__(
14391427 model : str ,
14401428 openai_client : AsyncOpenAI ,
14411429 ) -> None :
1442- global response_index
1443- response_index = 0
1444- global guardrail_response_index
1445- guardrail_response_index = 0
14461430 super ().__init__ (model , openai_client )
1431+ self ._responses = iter (self .responses )
1432+ self ._guardrail_responses = iter (self .guardrail_responses )
14471433
14481434 async def get_response (
14491435 self ,
@@ -1461,15 +1447,9 @@ async def get_response(
14611447 system_instructions
14621448 == "Check if the user is asking you to do their math homework."
14631449 ):
1464- global guardrail_response_index
1465- response = self .guardrail_responses [guardrail_response_index ]
1466- guardrail_response_index += 1
1467- return response
1450+ return next (self ._guardrail_responses )
14681451 else :
1469- global response_index
1470- response = self .responses [response_index ]
1471- response_index += 1
1472- return response
1452+ return next (self ._responses )
14731453
14741454
14751455### 1. An agent-based guardrail that is triggered if the user is asking to do math homework
0 commit comments