|
39 | 39 | # The project name is included in the CICD pipeline
|
40 | 40 | # os.environ['GOOGLE_CLOUD_PROJECT'] = "add-your-project-name"
|
41 | 41 |
|
| 42 | + |
42 | 43 | @pytest.fixture()
|
43 | 44 | def mock_rag_components(mocker):
|
44 |
| - mock_client_cls = mocker.patch("google.genai.Client") |
45 |
| - |
46 |
| - |
47 |
| - class AsyncIterator: |
48 |
| - def __aiter__(self): |
49 |
| - return self |
50 |
| - |
51 |
| - |
52 |
| - async def __anext__(self): |
53 |
| - if not hasattr(self, "used"): |
54 |
| - self.used = True |
55 |
| - return mocker.MagicMock( |
56 |
| - text="Mariusz Pudzianowski won in 2002, 2003, 2005, 2007, and 2008." |
57 |
| - ) |
58 |
| - raise StopAsyncIteration |
| 45 | + mock_client_cls = mocker.patch("google.genai.Client") |
59 | 46 |
|
| 47 | + class AsyncIterator: |
| 48 | + def __aiter__(self): |
| 49 | + return self |
60 | 50 |
|
61 |
| - mock_session = mocker.AsyncMock() |
62 |
| - mock_session.__aenter__.return_value = mock_session |
63 |
| - mock_session.receive = lambda: AsyncIterator() |
| 51 | + async def __anext__(self): |
| 52 | + if not hasattr(self, "used"): |
| 53 | + self.used = True |
| 54 | + return mocker.MagicMock( |
| 55 | + text="Mariusz Pudzianowski won in 2002, 2003, 2005, 2007, and 2008." |
| 56 | + ) |
| 57 | + raise StopAsyncIteration |
64 | 58 |
|
| 59 | + mock_session = mocker.AsyncMock() |
| 60 | + mock_session.__aenter__.return_value = mock_session |
| 61 | + mock_session.receive = lambda: AsyncIterator() |
65 | 62 |
|
66 |
| - mock_client_cls.return_value.aio.live.connect.return_value = mock_session |
| 63 | + mock_client_cls.return_value.aio.live.connect.return_value = mock_session |
67 | 64 |
|
68 | 65 |
|
69 | 66 | @pytest.mark.asyncio
|
@@ -129,4 +126,4 @@ async def test_live_structured_ouput_with_txt() -> None:
|
129 | 126 |
|
130 | 127 | @pytest.mark.asyncio
|
131 | 128 | async def test_live_ground_ragengine_with_txt(mock_rag_components) -> None:
|
132 |
| - assert await live_ground_ragengine_with_txt.generate_content("test") |
| 129 | + assert await live_ground_ragengine_with_txt.generate_content("test") |
0 commit comments