Repair tests for aiobotocore

Author: TimPansino

newrelic/hooks/external_botocore.py

@@ -877,43 +877,7 @@ def extract_bedrock_converse_attrs(kwargs, response, response_headers, model, sp
     return bedrock_attrs
 
 
-class EventStreamWrapper(ObjectProxy):
-    def __iter__(self):
-        g = GeneratorProxy(self.__wrapped__.__iter__())
-        g._nr_ft = getattr(self, "_nr_ft", None)
-        g._nr_bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
-        g._nr_model_extractor = getattr(self, "_nr_model_extractor", NULL_EXTRACTOR)
-        g._nr_is_converse = getattr(self, "_nr_is_converse", False)
-        return g
-
-
-class GeneratorProxy(ObjectProxy):
-    def __init__(self, wrapped):
-        super().__init__(wrapped)
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        transaction = current_transaction()
-        if not transaction:
-            return self.__wrapped__.__next__()
-
-        return_val = None
-        try:
-            return_val = self.__wrapped__.__next__()
-            self.record_stream_chunk(return_val, transaction)
-        except StopIteration:
-            self.record_events_on_stop_iteration(transaction)
-            raise
-        except Exception as exc:
-            self.record_error(transaction, exc)
-            raise
-        return return_val
-
-    def close(self):
-        return super().close()
-
+class BedrockRecordEventMixin:
     def record_events_on_stop_iteration(self, transaction):
         if hasattr(self, "_nr_ft"):
             bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
@@ -1002,6 +966,44 @@ def converse_record_stream_chunk(self, event, transaction):
         #     self.record_events_on_stop_iteration(transaction)
 
 
+class EventStreamWrapper(ObjectProxy):
+    def __iter__(self):
+        g = GeneratorProxy(self.__wrapped__.__iter__())
+        g._nr_ft = getattr(self, "_nr_ft", None)
+        g._nr_bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
+        g._nr_model_extractor = getattr(self, "_nr_model_extractor", NULL_EXTRACTOR)
+        g._nr_is_converse = getattr(self, "_nr_is_converse", False)
+        return g
+
+
+class GeneratorProxy(BedrockRecordEventMixin, ObjectProxy):
+    def __init__(self, wrapped):
+        super().__init__(wrapped)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        transaction = current_transaction()
+        if not transaction:
+            return self.__wrapped__.__next__()
+
+        return_val = None
+        try:
+            return_val = self.__wrapped__.__next__()
+            self.record_stream_chunk(return_val, transaction)
+        except StopIteration:
+            self.record_events_on_stop_iteration(transaction)
+            raise
+        except Exception as exc:
+            self.record_error(transaction, exc)
+            raise
+        return return_val
+
+    def close(self):
+        return super().close()
+
+
 class AsyncEventStreamWrapper(ObjectProxy):
     def __aiter__(self):
         g = AsyncGeneratorProxy(self.__wrapped__.__aiter__())
@@ -1012,13 +1014,7 @@ def __aiter__(self):
         return g
 
 
-class AsyncGeneratorProxy(ObjectProxy):
-    # Import these methods from the synchronous GeneratorProxy
-    # Avoid direct inheritance so we don't implement both __iter__ and __aiter__
-    record_stream_chunk = GeneratorProxy.record_stream_chunk
-    record_events_on_stop_iteration = GeneratorProxy.record_events_on_stop_iteration
-    record_error = GeneratorProxy.record_error
-
+class AsyncGeneratorProxy(BedrockRecordEventMixin, ObjectProxy):
     def __aiter__(self):
         return self
 

tests/external_aiobotocore/test_bedrock_chat_completion_invoke_model.py

@@ -20,6 +20,18 @@
 import botocore.exceptions
 import pytest
 from conftest import BOTOCORE_VERSION
+from external_botocore._test_bedrock_chat_completion_invoke_model import (
+    chat_completion_expected_events,
+    chat_completion_expected_malformed_request_body_events,
+    chat_completion_expected_malformed_response_body_events,
+    chat_completion_expected_malformed_response_streaming_body_events,
+    chat_completion_expected_malformed_response_streaming_chunk_events,
+    chat_completion_expected_streaming_error_events,
+    chat_completion_invalid_access_key_error_events,
+    chat_completion_invalid_model_error_events,
+    chat_completion_payload_templates,
+    chat_completion_streaming_expected_events,
+)
 from testing_support.fixtures import override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes
 from testing_support.ml_testing_utils import (
     add_token_count_to_events,
@@ -42,18 +54,6 @@
 from newrelic.api.transaction import add_custom_attribute
 from newrelic.common.object_names import callable_name
 from newrelic.hooks.external_botocore import MODEL_EXTRACTORS
-from tests.external_botocore._test_bedrock_chat_completion_invoke_model import (
-    chat_completion_expected_events,
-    chat_completion_expected_malformed_request_body_events,
-    chat_completion_expected_malformed_response_body_events,
-    chat_completion_expected_malformed_response_streaming_body_events,
-    chat_completion_expected_malformed_response_streaming_chunk_events,
-    chat_completion_expected_streaming_error_events,
-    chat_completion_invalid_access_key_error_events,
-    chat_completion_invalid_model_error_events,
-    chat_completion_payload_templates,
-    chat_completion_streaming_expected_events,
-)
 
 
 @pytest.fixture(scope="session", params=[False, True], ids=["ResponseStandard", "ResponseStreaming"])