Skip to content

Commit 5129b6f

Browse files
committed
Implement streaming support in RequestsMocker, get tests passing again.
1 parent 51bd151 commit 5129b6f

File tree

1 file changed

+10
-3
lines changed
  • instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/common

1 file changed

+10
-3
lines changed

instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/common/requests_mocker.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,13 @@ def combined_generator(args):
140140
response = generator(args)
141141
if response.status_code != 200:
142142
continue
143-
# TODO: ...
143+
response_json = response.json()
144+
response_json_str = json.dumps(response_json_str)
145+
contents.append(f'data: {response_json_str}')
146+
contents_str = '\r\n'.join(contents)
147+
full_contents = f'{contents_str}\r\n\r\n'
148+
result.raw = io.StringIO(full_contents)
149+
return result
144150
return combined_generator
145151

146152

@@ -183,8 +189,8 @@ def _do_send(
183189
):
184190
stream=kwargs.get('stream', False)
185191
if not stream:
186-
return _do_send_non_streaming(session, request, **kwargs)
187-
return _do_send_streaming(session, request, **kwargs)
192+
return self._do_send_non_streaming(session, request, **kwargs)
193+
return self._do_send_streaming(session, request, **kwargs)
188194

189195
def _do_send_streaming(
190196
self,
@@ -202,6 +208,7 @@ def _do_send_streaming(
202208
response_generators.append(_return_404)
203209
args = RequestsCallArgs(session, request, **kwargs)
204210
response_generator = _to_stream_response_generator(response_generators)
211+
call = RequestsCall(args, response_generator)
205212
result = call.response
206213
self._calls.append(call)
207214
return result

0 commit comments

Comments
 (0)