Coverage for tests/openaiclient_mock.py: 81%

1""" 

2# OpenAI client mock. 

3""" 

4 

5from unittest.mock import MagicMock 

6from unittest.mock import AsyncMock 

7 

8from typing import Dict 

9from typing import Any 

10 

11mock_openai = MagicMock() 

12mock_client = MagicMock() 

13mock_client.chat.completions.create = AsyncMock() 

14mock_openai.AsyncOpenAI.return_value = mock_client 

15 

16 

17class OpenAIClientMock: 

18 

19 def __init__(self) -> None: 

20 self.response = MagicMock() 

21 self.response.choices = [ 

22 MagicMock(message=MagicMock(content="mock response")) 

23 ] 

24 self.client = MagicMock() 

25 self.client.chat.completions.create = AsyncMock( 

26 return_value=self.response 

27 ) 

28 

29 self.AsyncOpenAI = MagicMock(return_value=self.client) 

30 self.AsyncAzureOpenAI = MagicMock(return_value=self.client) 

31 

32 def get_sub_modules(self) -> Dict[str, Any]: 

33 openai_module = MagicMock() 

34 openai_module.AsyncOpenAI = self.AsyncOpenAI 

35 openai_module.AsyncAzureOpenAI = self.AsyncAzureOpenAI 

36 return { 

37 "openai": openai_module, 

38 }