diff --git a/newrelic/hooks/datastore_pymongo.py b/newrelic/hooks/datastore_pymongo.py index e0593ed34..af5f6bf9b 100644 --- a/newrelic/hooks/datastore_pymongo.py +++ b/newrelic/hooks/datastore_pymongo.py @@ -18,7 +18,6 @@ from newrelic.api.function_trace import wrap_function_trace from newrelic.common.object_wrapper import wrap_function_wrapper - _pymongo_client_async_methods = ( "aggregate", "aggregate_raw_batches", @@ -138,6 +137,7 @@ def instrument_pymongo_synchronous_pool(module): module, "Connection.__init__", name=f"{module.__name__}:Connection.__init__", terminal=True, rollup=rollup ) + def instrument_pymongo_asynchronous_pool(module): rollup = ("Datastore/all", "Datastore/MongoDB/all") @@ -145,7 +145,11 @@ def instrument_pymongo_asynchronous_pool(module): # __getattr__() method in a way that breaks introspection. wrap_function_trace( - module, "AsyncConnection.__init__", name=f"{module.__name__}:AsyncConnection.__init__", terminal=True, rollup=rollup + module, + "AsyncConnection.__init__", + name=f"{module.__name__}:AsyncConnection.__init__", + terminal=True, + rollup=rollup, ) @@ -164,6 +168,7 @@ def instrument_pymongo_synchronous_mongo_client(module): module, "MongoClient.__init__", name=f"{module.__name__}:MongoClient.__init__", terminal=True, rollup=rollup ) + def instrument_pymongo_asynchronous_mongo_client(module): rollup = ("Datastore/all", "Datastore/MongoDB/all") @@ -171,7 +176,11 @@ def instrument_pymongo_asynchronous_mongo_client(module): # __getattr__() method in a way that breaks introspection. wrap_function_trace( - module, "AsyncMongoClient.__init__", name=f"{module.__name__}:AsyncMongoClient.__init__", terminal=True, rollup=rollup + module, + "AsyncMongoClient.__init__", + name=f"{module.__name__}:AsyncMongoClient.__init__", + terminal=True, + rollup=rollup, ) diff --git a/tests/datastore_pymongo/test_async_collection.py b/tests/datastore_pymongo/test_async_collection.py index b0e152aa4..9a1101b1c 100644 --- a/tests/datastore_pymongo/test_async_collection.py +++ b/tests/datastore_pymongo/test_async_collection.py @@ -14,15 +14,14 @@ import sqlite3 -import pytest - import pymongo +import pytest from pymongo.asynchronous.mongo_client import AsyncMongoClient from testing_support.db_settings import mongodb_settings -from testing_support.validators.validate_span_events import validate_span_events from testing_support.validators.validate_database_duration import ( validate_database_duration, ) +from testing_support.validators.validate_span_events import validate_span_events from testing_support.validators.validate_transaction_errors import ( validate_transaction_errors, ) @@ -34,7 +33,7 @@ from newrelic.common import system_info from newrelic.common.package_version_utils import get_package_version_tuple -# Skip if AsyncMongoClient has not been implemented yet +# Skip if AsyncMongoClient has not been implemented yet if get_package_version_tuple("pymongo") < (4, 9, 0): pytest.skip(allow_module_level=True, reason="AsyncMongoClient not available.") diff --git a/tests/datastore_pymongo/test_collection.py b/tests/datastore_pymongo/test_collection.py index ead723c41..ccf92e6b8 100644 --- a/tests/datastore_pymongo/test_collection.py +++ b/tests/datastore_pymongo/test_collection.py @@ -16,10 +16,10 @@ import pymongo from testing_support.db_settings import mongodb_settings -from testing_support.validators.validate_span_events import validate_span_events from testing_support.validators.validate_database_duration import ( validate_database_duration, ) +from testing_support.validators.validate_span_events import validate_span_events from testing_support.validators.validate_transaction_errors import ( validate_transaction_errors, ) diff --git a/tests/datastore_pymongo/test_uninstrumented_methods.py b/tests/datastore_pymongo/test_uninstrumented_methods.py index 2a6213347..ff4365c42 100644 --- a/tests/datastore_pymongo/test_uninstrumented_methods.py +++ b/tests/datastore_pymongo/test_uninstrumented_methods.py @@ -13,9 +13,8 @@ # limitations under the License. import pytest - -from pymongo.synchronous.mongo_client import MongoClient from pymongo.asynchronous.mongo_client import AsyncMongoClient +from pymongo.synchronous.mongo_client import MongoClient from testing_support.db_settings import mongodb_settings DB_SETTINGS = mongodb_settings()[0]