21
21
from contextlib import contextmanager
22
22
from typing import Any
23
23
from unittest import TestCase
24
- from unittest .mock import MagicMock , patch
24
+ from unittest .mock import ANY , MagicMock , patch
25
25
26
26
import fsspec
27
27
@@ -119,8 +119,9 @@ def setUp(self):
119
119
self ._fsspec_patcher = patch (
120
120
"opentelemetry.util.genai._fsspec_upload.completion_hook.fsspec"
121
121
)
122
- self .mock_fsspec = self ._fsspec_patcher .start ()
123
- self .mock_fsspec .open = ThreadSafeMagicMock ()
122
+ mock_fsspec = self ._fsspec_patcher .start ()
123
+ self .mock_fs = ThreadSafeMagicMock ()
124
+ mock_fsspec .url_to_fs .return_value = self .mock_fs , ""
124
125
125
126
self .hook = FsspecUploadCompletionHook (
126
127
base_path = BASE_PATH ,
@@ -135,12 +136,12 @@ def tearDown(self) -> None:
135
136
def block_upload (self ):
136
137
unblock_upload = threading .Event ()
137
138
138
- def blocked_upload (* args : Any ):
139
+ def blocked_upload (* args : Any , ** kwargs : Any ):
139
140
unblock_upload .wait ()
140
141
return MagicMock ()
141
142
142
143
try :
143
- self .mock_fsspec .open .side_effect = blocked_upload
144
+ self .mock_fs .open .side_effect = blocked_upload
144
145
yield
145
146
finally :
146
147
unblock_upload .set ()
@@ -158,7 +159,7 @@ def test_upload_then_shutdown(self):
158
159
self .hook .shutdown ()
159
160
160
161
self .assertEqual (
161
- self .mock_fsspec .open .call_count ,
162
+ self .mock_fs .open .call_count ,
162
163
3 ,
163
164
"should have uploaded 3 files" ,
164
165
)
@@ -174,7 +175,7 @@ def test_upload_blocked(self):
174
175
)
175
176
176
177
self .assertLessEqual (
177
- self .mock_fsspec .open .call_count ,
178
+ self .mock_fs .open .call_count ,
178
179
MAXSIZE ,
179
180
f"uploader should only be called { MAXSIZE = } times" ,
180
181
)
@@ -202,7 +203,7 @@ def test_shutdown_timeout(self):
202
203
self .hook .shutdown (timeout_sec = 0.01 )
203
204
204
205
def test_failed_upload_logs (self ):
205
- self .mock_fsspec .open .side_effect = RuntimeError ("failed to upload" )
206
+ self .mock_fs .open .side_effect = RuntimeError ("failed to upload" )
206
207
207
208
with self .assertLogs (level = logging .ERROR ) as logs :
208
209
self .hook .on_completion (
@@ -220,6 +221,27 @@ def test_invalid_upload_format(self):
220
221
base_path = BASE_PATH , upload_format = "invalid"
221
222
)
222
223
224
+ def test_upload_format_sets_content_type (self ):
225
+ for upload_format , expect_content_type in (
226
+ ("json" , "application/json" ),
227
+ ("jsonl" , "application/jsonl" ),
228
+ ):
229
+ hook = FsspecUploadCompletionHook (
230
+ base_path = BASE_PATH , upload_format = upload_format
231
+ )
232
+ self .addCleanup (hook .shutdown )
233
+
234
+ hook .on_completion (
235
+ inputs = FAKE_INPUTS ,
236
+ outputs = FAKE_OUTPUTS ,
237
+ system_instruction = FAKE_SYSTEM_INSTRUCTION ,
238
+ )
239
+ hook .shutdown ()
240
+
241
+ self .mock_fs .open .assert_called_with (
242
+ ANY , "w" , content_type = expect_content_type
243
+ )
244
+
223
245
def test_parse_upload_format_envvar (self ):
224
246
for envvar_value , expect in (
225
247
("" , "json" ),
@@ -250,7 +272,11 @@ def test_parse_upload_format_envvar(self):
250
272
base_path = BASE_PATH , upload_format = "jsonl"
251
273
)
252
274
self .addCleanup (hook .shutdown )
253
- self .assertEqual (hook ._format , "jsonl" )
275
+ self .assertEqual (
276
+ hook ._format ,
277
+ "jsonl" ,
278
+ "upload_format kwarg should take precedence" ,
279
+ )
254
280
255
281
def test_upload_after_shutdown_logs (self ):
256
282
self .hook .shutdown ()
@@ -417,3 +443,26 @@ def test_upload_jsonlines(self) -> None:
417
443
{"role":"user","parts":[{"response":{"capital":"Paris"},"id":"get_capital_0","type":"tool_call_response"}],"index":2}
418
444
""" ,
419
445
)
446
+
447
+ def test_upload_chained_filesystem_ref (self ) -> None :
448
+ """Using a chained filesystem like simplecache should refer to the final remote destination"""
449
+ hook = FsspecUploadCompletionHook (
450
+ base_path = "simplecache::memory" ,
451
+ upload_format = "jsonl" ,
452
+ )
453
+ self .addCleanup (hook .shutdown )
454
+ log_record = LogRecord ()
455
+
456
+ hook .on_completion (
457
+ inputs = FAKE_INPUTS ,
458
+ outputs = FAKE_OUTPUTS ,
459
+ system_instruction = FAKE_SYSTEM_INSTRUCTION ,
460
+ log_record = log_record ,
461
+ )
462
+ hook .shutdown ()
463
+
464
+ ref_uri : str = log_record .attributes ["gen_ai.input.messages_ref" ]
465
+ self .assertTrue (
466
+ ref_uri .startswith ("memory://" ),
467
+ f"{ ref_uri = } does not start with final destination uri memory://" ,
468
+ )
0 commit comments