Add missing content formatter error test case

Author: amenegola

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

@@ -125,8 +125,7 @@ def mock_bq_client():
 def mock_write_client():
   # Updated patch path to match the new import structure in src
   with mock.patch(
-      "google.cloud.bigquery_storage_v1.services.big_query_write.async_client.BigQueryWriteAsyncClient",
-      autospec=True,
+      "google.cloud.bigquery_storage_v1.BigQueryWriteAsyncClient", autospec=True
   ) as mock_cls:
     mock_client = mock_cls.return_value
     mock_client.transport = mock.AsyncMock()
@@ -298,13 +297,7 @@ async def test_event_allowlist(
     mock_write_client.append_rows.assert_called_once()
     mock_write_client.append_rows.reset_mock()
 
-    # Re-init plugin logic since close() shuts it down, but for this test we want to test denial
-    # However, close() cleans up clients. We should probably create a new plugin or just check that the task was not created.
-    # But on_user_message_callback will try to log.
-    # To keep it simple, let's just use a fresh plugin for the second part or assume close() resets state enough to re-run _ensure_init if needed,
-    # but _ensure_init is called inside _perform_write.
-    # Actually, close() sets _is_shutting_down to True, so further logs are ignored.
-    # So we need a new plugin instance or reset _is_shutting_down.
+    # Re-init plugin logic since close() shuts it down
     plugin._is_shutting_down = False
 
     user_message = types.Content(parts=[types.Part(text="What is up?")])
@@ -399,6 +392,44 @@ def mutate_payload(data):
     assert content["model"] == "GEMINI-PRO"
     assert content["prompt"][0]["role"] == "user"
 
+  @pytest.mark.asyncio
+  async def test_content_formatter_error_fallback(
+      self,
+      mock_write_client,
+      invocation_context,
+      mock_auth_default,
+      mock_bq_client,
+      mock_to_arrow_schema,
+      dummy_arrow_schema,
+      mock_asyncio_to_thread,
+  ):
+    """Tests that if content_formatter fails, the original payload is used."""
+
+    def error_formatter(data):
+      raise ValueError("Formatter failed")
+
+    config = BigQueryLoggerConfig(content_formatter=error_formatter)
+    plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
+        PROJECT_ID, DATASET_ID, TABLE_ID, config
+    )
+    await plugin._ensure_init()
+    mock_write_client.append_rows.reset_mock()
+
+    user_message = types.Content(parts=[types.Part(text="Original message")])
+
+    # This triggers the log. Internal logic catches exception and proceeds.
+    await plugin.on_user_message_callback(
+        invocation_context=invocation_context, user_message=user_message
+    )
+    await plugin.close()
+
+    mock_write_client.append_rows.assert_called_once()
+    log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
+
+    # Verify that despite the error, we still got the original data
+    content = json.loads(log_entry["content"])
+    assert content["text"] == "Original message"
+
   @pytest.mark.asyncio
   async def test_max_content_length_smart_truncation(
       self,
@@ -695,7 +726,9 @@ async def test_after_model_callback_tool_call(
       callback_context,
       dummy_arrow_schema,
   ):
-    tool_fc = types.FunctionCall(name="get_weather", args={"location": "Paris"})
+    tool_fc = types.FunctionCall(
+        name="get_weather", args={"location": "Paris"}
+    )
     llm_response = llm_response_lib.LlmResponse(
         content=types.Content(parts=[types.Part(function_call=tool_fc)]),
         usage_metadata=types.UsageMetadata(
@@ -725,7 +758,9 @@ async def test_before_tool_callback_logs_correctly(
     type(mock_tool).name = mock.PropertyMock(return_value="MyTool")
     type(mock_tool).description = mock.PropertyMock(return_value="Description")
     await bq_plugin_inst.before_tool_callback(
-        tool=mock_tool, tool_args={"param": "value"}, tool_context=tool_context
+        tool=mock_tool,
+        tool_args={"param": "value"},
+        tool_context=tool_context,
     )
     await bq_plugin_inst.close()
     log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
@@ -894,7 +929,9 @@ async def fake_append_rows_with_schema_error(requests, **kwargs):
       )
 
   @pytest.mark.asyncio
-  async def test_close(self, bq_plugin_inst, mock_bq_client, mock_write_client):
+  async def test_close(
+      self, bq_plugin_inst, mock_bq_client, mock_write_client
+  ):
     await bq_plugin_inst.close()
     mock_write_client.transport.close.assert_called_once()
     # in the new implementation we verify attributes are reset
@@ -997,4 +1034,4 @@ async def test_table_creation_options(
 
     # Verify schema type for content is JSON
     content_field = next(f for f in table_arg.schema if f.name == "content")
-    assert content_field.field_type == "JSON"
+    assert content_field.field_type == "JSON"