| 42 | |
| 43 | |
| 44 | class MockOpenAIClient: |
| 45 | def __init__(self): |
| 46 | self.chat = MagicMock() |
| 47 | self.chat.completions = MagicMock() |
| 48 | |
| 49 | def set_response(self, response: ChatCompletion): |
| 50 | """ |
| 51 | Set the mock to return a specific response. |
| 52 | :param response: A ChatCompletion response to return. |
| 53 | """ |
| 54 | self.chat.completions.create.return_value = response |
| 55 | |
| 56 | def set_sequential_responses(self, responses: list[ChatCompletion]): |
| 57 | """ |
| 58 | Set the mock to return different responses sequentially. |
| 59 | :param responses: A list of ChatCompletion responses to return in order. |
| 60 | """ |
| 61 | self.chat.completions.create.side_effect = responses |
| 62 | |
| 63 | def assert_create_called_with(self, **kwargs): |
| 64 | self.chat.completions.create.assert_called_with(**kwargs) |
| 65 | |
| 66 | |
| 67 | # Initialize the mock client |
no outgoing calls
searching dependent graphs…