|
5 | 5 |
|
6 | 6 | import time
|
7 | 7 | import unittest
|
8 |
| -from unittest.mock import Mock, patch |
| 8 | +from unittest.mock import MagicMock, Mock, patch |
9 | 9 |
|
10 | 10 | from botocore.exceptions import ClientError
|
11 | 11 |
|
@@ -579,6 +579,134 @@ def test_is_batch_active_edge_cases(self):
|
579 | 579 | result = self.log_client._is_batch_active(batch2, current_time)
|
580 | 580 | self.assertFalse(result)
|
581 | 581 |
|
| 582 | + @patch("amazon.opentelemetry.distro.exporter.aws.metrics._cloudwatch_log_client.suppress_instrumentation") |
| 583 | + def test_create_log_group_uses_suppress_instrumentation(self, mock_suppress): |
| 584 | + """Test that _create_log_group_if_needed uses suppress_instrumentation.""" |
| 585 | + # Configure the mock context manager |
| 586 | + mock_context = MagicMock() |
| 587 | + mock_suppress.return_value = mock_context |
| 588 | + mock_context.__enter__.return_value = mock_context |
| 589 | + mock_context.__exit__.return_value = None |
| 590 | + |
| 591 | + # Call the method |
| 592 | + self.log_client._create_log_group_if_needed() |
| 593 | + |
| 594 | + # Verify suppress_instrumentation was called |
| 595 | + mock_suppress.assert_called_once() |
| 596 | + mock_context.__enter__.assert_called_once() |
| 597 | + mock_context.__exit__.assert_called_once() |
| 598 | + |
| 599 | + # Verify the AWS call happened within the context |
| 600 | + self.log_client.logs_client.create_log_group.assert_called_once_with(logGroupName="test-log-group") |
| 601 | + |
| 602 | + @patch("amazon.opentelemetry.distro.exporter.aws.metrics._cloudwatch_log_client.suppress_instrumentation") |
| 603 | + def test_create_log_stream_uses_suppress_instrumentation(self, mock_suppress): |
| 604 | + """Test that _create_log_stream_if_needed uses suppress_instrumentation.""" |
| 605 | + # Configure the mock context manager |
| 606 | + mock_context = MagicMock() |
| 607 | + mock_suppress.return_value = mock_context |
| 608 | + mock_context.__enter__.return_value = mock_context |
| 609 | + mock_context.__exit__.return_value = None |
| 610 | + |
| 611 | + # Call the method |
| 612 | + self.log_client._create_log_stream_if_needed() |
| 613 | + |
| 614 | + # Verify suppress_instrumentation was called |
| 615 | + mock_suppress.assert_called_once() |
| 616 | + mock_context.__enter__.assert_called_once() |
| 617 | + mock_context.__exit__.assert_called_once() |
| 618 | + |
| 619 | + # Verify the AWS call happened within the context |
| 620 | + self.log_client.logs_client.create_log_stream.assert_called_once() |
| 621 | + |
| 622 | + @patch("amazon.opentelemetry.distro.exporter.aws.metrics._cloudwatch_log_client.suppress_instrumentation") |
| 623 | + def test_send_log_batch_uses_suppress_instrumentation(self, mock_suppress): |
| 624 | + """Test that _send_log_batch uses suppress_instrumentation.""" |
| 625 | + # Configure the mock context manager |
| 626 | + mock_context = MagicMock() |
| 627 | + mock_suppress.return_value = mock_context |
| 628 | + mock_context.__enter__.return_value = mock_context |
| 629 | + mock_context.__exit__.return_value = None |
| 630 | + |
| 631 | + # Create a batch with events |
| 632 | + batch = self.log_client._create_event_batch() |
| 633 | + batch.add_event({"message": "test", "timestamp": int(time.time() * 1000)}, 10) |
| 634 | + |
| 635 | + # Mock successful put_log_events |
| 636 | + self.log_client.logs_client.put_log_events.return_value = {"nextSequenceToken": "12345"} |
| 637 | + |
| 638 | + # Call the method |
| 639 | + self.log_client._send_log_batch(batch) |
| 640 | + |
| 641 | + # Verify suppress_instrumentation was called |
| 642 | + mock_suppress.assert_called_once() |
| 643 | + mock_context.__enter__.assert_called_once() |
| 644 | + mock_context.__exit__.assert_called_once() |
| 645 | + |
| 646 | + # Verify the AWS call happened within the context |
| 647 | + self.log_client.logs_client.put_log_events.assert_called_once() |
| 648 | + |
| 649 | + @patch("amazon.opentelemetry.distro.exporter.aws.metrics._cloudwatch_log_client.suppress_instrumentation") |
| 650 | + def test_send_log_batch_retry_uses_suppress_instrumentation(self, mock_suppress): |
| 651 | + """Test that _send_log_batch retry logic also uses suppress_instrumentation.""" |
| 652 | + # Configure the mock context manager |
| 653 | + mock_context = MagicMock() |
| 654 | + mock_suppress.return_value = mock_context |
| 655 | + mock_context.__enter__.return_value = mock_context |
| 656 | + mock_context.__exit__.return_value = None |
| 657 | + |
| 658 | + # Create a batch with events |
| 659 | + batch = self.log_client._create_event_batch() |
| 660 | + batch.add_event({"message": "test", "timestamp": int(time.time() * 1000)}, 10) |
| 661 | + |
| 662 | + # Mock put_log_events to fail first with ResourceNotFoundException, then succeed |
| 663 | + self.log_client.logs_client.put_log_events.side_effect = [ |
| 664 | + ClientError({"Error": {"Code": "ResourceNotFoundException"}}, "PutLogEvents"), |
| 665 | + {"nextSequenceToken": "12345"}, |
| 666 | + ] |
| 667 | + |
| 668 | + # Call the method |
| 669 | + self.log_client._send_log_batch(batch) |
| 670 | + |
| 671 | + # Verify suppress_instrumentation was called: |
| 672 | + # 1. Initial _send_log_batch context |
| 673 | + # 2. Nested context in the retry block |
| 674 | + # 3. _create_log_group_if_needed context |
| 675 | + # 4. _create_log_stream_if_needed context |
| 676 | + self.assertEqual(mock_suppress.call_count, 4) |
| 677 | + # Each context should have been properly entered and exited |
| 678 | + self.assertEqual(mock_context.__enter__.call_count, 4) |
| 679 | + self.assertEqual(mock_context.__exit__.call_count, 4) |
| 680 | + |
| 681 | + # Verify AWS calls happened |
| 682 | + self.assertEqual(self.log_client.logs_client.put_log_events.call_count, 2) |
| 683 | + self.log_client.logs_client.create_log_group.assert_called_once() |
| 684 | + self.log_client.logs_client.create_log_stream.assert_called_once() |
| 685 | + |
| 686 | + @patch("amazon.opentelemetry.distro.exporter.aws.metrics._cloudwatch_log_client.suppress_instrumentation") |
| 687 | + def test_create_log_group_exception_still_uses_suppress_instrumentation(self, mock_suppress): |
| 688 | + """Test that suppress_instrumentation is properly used even when exceptions occur.""" |
| 689 | + # Configure the mock context manager |
| 690 | + mock_context = MagicMock() |
| 691 | + mock_suppress.return_value = mock_context |
| 692 | + mock_context.__enter__.return_value = mock_context |
| 693 | + mock_context.__exit__.return_value = None |
| 694 | + |
| 695 | + # Make create_log_group raise an exception |
| 696 | + self.log_client.logs_client.create_log_group.side_effect = ClientError( |
| 697 | + {"Error": {"Code": "AccessDenied"}}, "CreateLogGroup" |
| 698 | + ) |
| 699 | + |
| 700 | + # Call should raise the exception |
| 701 | + with self.assertRaises(ClientError): |
| 702 | + self.log_client._create_log_group_if_needed() |
| 703 | + |
| 704 | + # Verify suppress_instrumentation was still properly used |
| 705 | + mock_suppress.assert_called_once() |
| 706 | + mock_context.__enter__.assert_called_once() |
| 707 | + # __exit__ should be called even though an exception was raised |
| 708 | + mock_context.__exit__.assert_called_once() |
| 709 | + |
582 | 710 |
|
583 | 711 | if __name__ == "__main__":
|
584 | 712 | unittest.main()
|
0 commit comments