def test_wrap_connect(self):
aiopg_mock = AiopgMock()
with mock.patch("aiopg.connect", aiopg_mock.connect):
wrappers.wrap_connect(self.tracer, "-")
connection = async_call(aiopg.connect())
self.assertEqual(aiopg_mock.connect_call_count, 1)
self.assertIsInstance(connection.__wrapped__, mock.Mock)
def _instrument(self, **kwargs):
"""Integrate with PostgreSQL aiopg library.
aiopg: https://github.com/aio-libs/aiopg
"""
tracer_provider = kwargs.get("tracer_provider")
wrappers.wrap_connect(
__name__,
self._DATABASE_SYSTEM,
self._CONNECTION_ATTRIBUTES,
version=__version__,
tracer_provider=tracer_provider,
)
wrappers.wrap_create_pool(
__name__,
self._DATABASE_SYSTEM,
self._CONNECTION_ATTRIBUTES,
version=__version__,
tracer_provider=tracer_provider,
)
