|
6 | 6 | import pytest |
7 | 7 |
|
8 | 8 | from lightspeed_agent.api.a2a.logging_plugin import AgentLoggingPlugin, _truncate |
| 9 | +from lightspeed_agent.auth.middleware import ( |
| 10 | + _request_id, |
| 11 | + _request_order_id, |
| 12 | + _request_org_id, |
| 13 | + _request_user_id, |
| 14 | +) |
9 | 15 |
|
10 | 16 |
|
11 | 17 | class TestTruncate: |
@@ -331,3 +337,155 @@ async def test_all_callbacks_return_none(self, plugin): |
331 | 337 | ) |
332 | 338 | is None |
333 | 339 | ) |
| 340 | + |
| 341 | + |
| 342 | +class TestAuditFieldsInLogMessages: |
| 343 | + """Verify that audit context fields (user_id, org_id, order_id, request_id) |
| 344 | + are included in log messages when contextvars are set.""" |
| 345 | + |
| 346 | + @pytest.fixture |
| 347 | + def plugin(self): |
| 348 | + with patch( |
| 349 | + "lightspeed_agent.api.a2a.logging_plugin.get_settings" |
| 350 | + ) as mock_settings: |
| 351 | + settings = MagicMock() |
| 352 | + settings.agent_logging_detail = "basic" |
| 353 | + mock_settings.return_value = settings |
| 354 | + yield AgentLoggingPlugin() |
| 355 | + |
| 356 | + @pytest.fixture(autouse=True) |
| 357 | + def _set_audit_context(self): |
| 358 | + """Set audit contextvars for the duration of each test.""" |
| 359 | + token_user = _request_user_id.set("test-user-42") |
| 360 | + token_org = _request_org_id.set("test-org-7") |
| 361 | + token_order = _request_order_id.set("test-order-99") |
| 362 | + token_req = _request_id.set("req-abc-123") |
| 363 | + yield |
| 364 | + _request_user_id.reset(token_user) |
| 365 | + _request_org_id.reset(token_org) |
| 366 | + _request_order_id.reset(token_order) |
| 367 | + _request_id.reset(token_req) |
| 368 | + |
| 369 | + @pytest.mark.asyncio |
| 370 | + async def test_before_run_includes_audit_fields(self, plugin, caplog): |
| 371 | + ctx = MagicMock() |
| 372 | + ctx.invocation_id = "inv-1" |
| 373 | + ctx.agent = MagicMock(name="agent") |
| 374 | + |
| 375 | + with caplog.at_level(logging.INFO): |
| 376 | + await plugin.before_run_callback(invocation_context=ctx) |
| 377 | + |
| 378 | + assert "user_id=test-user-42" in caplog.text |
| 379 | + assert "org_id=test-org-7" in caplog.text |
| 380 | + assert "order_id=test-order-99" in caplog.text |
| 381 | + assert "request_id=req-abc-123" in caplog.text |
| 382 | + |
| 383 | + @pytest.mark.asyncio |
| 384 | + async def test_after_run_includes_audit_fields(self, plugin, caplog): |
| 385 | + ctx = MagicMock() |
| 386 | + ctx.invocation_id = "inv-1" |
| 387 | + |
| 388 | + with caplog.at_level(logging.INFO): |
| 389 | + await plugin.after_run_callback(invocation_context=ctx) |
| 390 | + |
| 391 | + assert "user_id=test-user-42" in caplog.text |
| 392 | + assert "org_id=test-org-7" in caplog.text |
| 393 | + assert "order_id=test-order-99" in caplog.text |
| 394 | + assert "request_id=req-abc-123" in caplog.text |
| 395 | + |
| 396 | + @pytest.mark.asyncio |
| 397 | + async def test_before_model_includes_audit_fields(self, plugin, caplog): |
| 398 | + ctx = MagicMock() |
| 399 | + ctx.agent_name = "agent" |
| 400 | + |
| 401 | + with caplog.at_level(logging.INFO): |
| 402 | + await plugin.before_model_callback( |
| 403 | + callback_context=ctx, llm_request=MagicMock() |
| 404 | + ) |
| 405 | + |
| 406 | + assert "user_id=test-user-42" in caplog.text |
| 407 | + assert "org_id=test-org-7" in caplog.text |
| 408 | + assert "order_id=test-order-99" in caplog.text |
| 409 | + assert "request_id=req-abc-123" in caplog.text |
| 410 | + |
| 411 | + @pytest.mark.asyncio |
| 412 | + async def test_after_model_includes_audit_fields(self, plugin, caplog): |
| 413 | + ctx = MagicMock() |
| 414 | + llm_response = MagicMock() |
| 415 | + llm_response.usage_metadata = None |
| 416 | + llm_response.model_version = None |
| 417 | + llm_response.content = None |
| 418 | + |
| 419 | + with caplog.at_level(logging.INFO): |
| 420 | + await plugin.after_model_callback( |
| 421 | + callback_context=ctx, llm_response=llm_response |
| 422 | + ) |
| 423 | + |
| 424 | + assert "user_id=test-user-42" in caplog.text |
| 425 | + assert "org_id=test-org-7" in caplog.text |
| 426 | + assert "order_id=test-order-99" in caplog.text |
| 427 | + assert "request_id=req-abc-123" in caplog.text |
| 428 | + |
| 429 | + @pytest.mark.asyncio |
| 430 | + async def test_before_tool_includes_audit_fields(self, plugin, caplog): |
| 431 | + tool = MagicMock() |
| 432 | + tool.name = "get_advisories" |
| 433 | + |
| 434 | + with caplog.at_level(logging.INFO): |
| 435 | + await plugin.before_tool_callback( |
| 436 | + tool=tool, tool_args={}, tool_context=MagicMock() |
| 437 | + ) |
| 438 | + |
| 439 | + assert "user_id=test-user-42" in caplog.text |
| 440 | + assert "org_id=test-org-7" in caplog.text |
| 441 | + assert "order_id=test-order-99" in caplog.text |
| 442 | + assert "request_id=req-abc-123" in caplog.text |
| 443 | + |
| 444 | + @pytest.mark.asyncio |
| 445 | + async def test_after_tool_includes_audit_fields(self, plugin, caplog): |
| 446 | + tool = MagicMock() |
| 447 | + tool.name = "get_advisories" |
| 448 | + |
| 449 | + with caplog.at_level(logging.INFO): |
| 450 | + await plugin.after_tool_callback( |
| 451 | + tool=tool, tool_args={}, tool_context=MagicMock(), result={} |
| 452 | + ) |
| 453 | + |
| 454 | + assert "user_id=test-user-42" in caplog.text |
| 455 | + assert "org_id=test-org-7" in caplog.text |
| 456 | + assert "order_id=test-order-99" in caplog.text |
| 457 | + assert "request_id=req-abc-123" in caplog.text |
| 458 | + |
| 459 | + @pytest.mark.asyncio |
| 460 | + async def test_tool_error_includes_audit_fields(self, plugin, caplog): |
| 461 | + tool = MagicMock() |
| 462 | + tool.name = "get_advisories" |
| 463 | + |
| 464 | + with caplog.at_level(logging.ERROR): |
| 465 | + await plugin.on_tool_error_callback( |
| 466 | + tool=tool, |
| 467 | + tool_args={}, |
| 468 | + tool_context=MagicMock(), |
| 469 | + error=RuntimeError("fail"), |
| 470 | + ) |
| 471 | + |
| 472 | + assert "user_id=test-user-42" in caplog.text |
| 473 | + assert "org_id=test-org-7" in caplog.text |
| 474 | + assert "order_id=test-order-99" in caplog.text |
| 475 | + assert "request_id=req-abc-123" in caplog.text |
| 476 | + |
| 477 | + @pytest.mark.asyncio |
| 478 | + async def test_model_error_includes_audit_fields(self, plugin, caplog): |
| 479 | + ctx = MagicMock() |
| 480 | + |
| 481 | + with caplog.at_level(logging.ERROR): |
| 482 | + await plugin.on_model_error_callback( |
| 483 | + callback_context=ctx, |
| 484 | + llm_request=MagicMock(), |
| 485 | + error=RuntimeError("fail"), |
| 486 | + ) |
| 487 | + |
| 488 | + assert "user_id=test-user-42" in caplog.text |
| 489 | + assert "org_id=test-org-7" in caplog.text |
| 490 | + assert "order_id=test-order-99" in caplog.text |
| 491 | + assert "request_id=req-abc-123" in caplog.text |
0 commit comments