@@ -303,3 +303,127 @@ def test_plugin_name_default(self):
303303 """Test that plugin has correct default name."""
304304 plugin = SaveFilesAsArtifactsPlugin ()
305305 assert plugin .name == "save_files_as_artifacts_plugin"
306+
307+ @pytest .mark .asyncio
308+ async def test_file_size_exceeds_limit (self ):
309+ """Test that files exceeding 20MB limit show error message."""
310+ # Create a file larger than 20MB (20 * 1024 * 1024 bytes)
311+ large_file_data = b"x" * (21 * 1024 * 1024 ) # 21 MB
312+ inline_data = types .Blob (
313+ display_name = "large_file.pdf" ,
314+ data = large_file_data ,
315+ mime_type = "application/pdf" ,
316+ )
317+
318+ user_message = types .Content (parts = [types .Part (inline_data = inline_data )])
319+
320+ result = await self .plugin .on_user_message_callback (
321+ invocation_context = self .mock_context , user_message = user_message
322+ )
323+
324+ # Should not try to save the artifact
325+ self .mock_context .artifact_service .save_artifact .assert_not_called ()
326+
327+ # Should return modified content with error message
328+ assert result is not None
329+ assert len (result .parts ) == 1
330+ assert result .parts [0 ].text is not None
331+ assert "[Upload Error:" in result .parts [0 ].text
332+ assert "21.00 MB" in result .parts [0 ].text
333+ assert "20 MB" in result .parts [0 ].text
334+ assert "Files API" in result .parts [0 ].text
335+
336+ @pytest .mark .asyncio
337+ async def test_file_size_at_limit (self ):
338+ """Test that files exactly at 20MB limit are processed successfully."""
339+ # Create a file exactly 20MB (20 * 1024 * 1024 bytes)
340+ file_data = b"x" * (20 * 1024 * 1024 ) # Exactly 20 MB
341+ inline_data = types .Blob (
342+ display_name = "max_size_file.pdf" ,
343+ data = file_data ,
344+ mime_type = "application/pdf" ,
345+ )
346+
347+ user_message = types .Content (parts = [types .Part (inline_data = inline_data )])
348+
349+ result = await self .plugin .on_user_message_callback (
350+ invocation_context = self .mock_context , user_message = user_message
351+ )
352+
353+ # Should save the artifact since it's at the limit
354+ self .mock_context .artifact_service .save_artifact .assert_called_once ()
355+ assert result is not None
356+ assert len (result .parts ) == 2
357+ assert result .parts [0 ].text == '[Uploaded Artifact: "max_size_file.pdf"]'
358+ assert result .parts [1 ].file_data is not None
359+
360+ @pytest .mark .asyncio
361+ async def test_file_size_just_over_limit (self ):
362+ """Test that files just over 20MB limit show error message."""
363+ # Create a file just over 20MB
364+ large_file_data = b"x" * (20 * 1024 * 1024 + 1 ) # 20 MB + 1 byte
365+ inline_data = types .Blob (
366+ display_name = "slightly_too_large.pdf" ,
367+ data = large_file_data ,
368+ mime_type = "application/pdf" ,
369+ )
370+
371+ user_message = types .Content (parts = [types .Part (inline_data = inline_data )])
372+
373+ result = await self .plugin .on_user_message_callback (
374+ invocation_context = self .mock_context , user_message = user_message
375+ )
376+
377+ # Should not try to save the artifact
378+ self .mock_context .artifact_service .save_artifact .assert_not_called ()
379+
380+ # Should return error message
381+ assert result is not None
382+ assert len (result .parts ) == 1
383+ assert "[Upload Error:" in result .parts [0 ].text
384+
385+ @pytest .mark .asyncio
386+ async def test_mixed_file_sizes (self ):
387+ """Test processing multiple files with mixed sizes."""
388+ # Small file (should succeed)
389+ small_file_data = b"x" * (5 * 1024 * 1024 ) # 5 MB
390+ small_inline_data = types .Blob (
391+ display_name = "small.pdf" ,
392+ data = small_file_data ,
393+ mime_type = "application/pdf" ,
394+ )
395+
396+ # Large file (should fail)
397+ large_file_data = b"x" * (25 * 1024 * 1024 ) # 25 MB
398+ large_inline_data = types .Blob (
399+ display_name = "large.pdf" ,
400+ data = large_file_data ,
401+ mime_type = "application/pdf" ,
402+ )
403+
404+ user_message = types .Content (
405+ parts = [
406+ types .Part (inline_data = small_inline_data ),
407+ types .Part (inline_data = large_inline_data ),
408+ ]
409+ )
410+
411+ result = await self .plugin .on_user_message_callback (
412+ invocation_context = self .mock_context , user_message = user_message
413+ )
414+
415+ # Should only save the small file
416+ self .mock_context .artifact_service .save_artifact .assert_called_once_with (
417+ app_name = "test_app" ,
418+ user_id = "test_user" ,
419+ session_id = "test_session" ,
420+ filename = "small.pdf" ,
421+ artifact = user_message .parts [0 ],
422+ )
423+
424+ # Should return both success and error messages
425+ assert result is not None
426+ assert len (result .parts ) == 3 # [success placeholder, file_data, error]
427+ assert '[Uploaded Artifact: "small.pdf"]' in result .parts [0 ].text
428+ assert result .parts [1 ].file_data is not None
429+ assert "[Upload Error:" in result .parts [2 ].text
0 commit comments