Refactor TracedConnectionProxy (#1097)

* Refactor TracedConnectionProxy

Fixes #1077

* Remove unecessary methods
This commit is contained in:
Diego Hurtado
2022-06-07 18:20:59 +01:00
committed by GitHub
parent 8afbce7533
commit 6876ad857f
4 changed files with 74 additions and 42 deletions

View File

@ -262,14 +262,14 @@ class TestDBApiIntegration(TestBase):
@mock.patch("opentelemetry.instrumentation.dbapi")
def test_wrap_connect(self, mock_dbapi):
dbapi.wrap_connect(self.tracer, mock_dbapi, "connect", "-")
dbapi.wrap_connect(self.tracer, MockConnectionEmpty(), "connect", "-")
connection = mock_dbapi.connect()
self.assertEqual(mock_dbapi.connect.call_count, 1)
self.assertIsInstance(connection.__wrapped__, mock.Mock)
self.assertIsInstance(connection._connection, mock.Mock)
@mock.patch("opentelemetry.instrumentation.dbapi")
def test_unwrap_connect(self, mock_dbapi):
dbapi.wrap_connect(self.tracer, mock_dbapi, "connect", "-")
dbapi.wrap_connect(self.tracer, MockConnectionEmpty(), "connect", "-")
connection = mock_dbapi.connect()
self.assertEqual(mock_dbapi.connect.call_count, 1)
@ -279,19 +279,21 @@ class TestDBApiIntegration(TestBase):
self.assertIsInstance(connection, mock.Mock)
def test_instrument_connection(self):
connection = mock.Mock()
connection = MockConnectionEmpty()
# Avoid get_attributes failing because can't concatenate mock
# pylint: disable=attribute-defined-outside-init
connection.database = "-"
connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
self.assertIs(connection2.__wrapped__, connection)
self.assertIs(connection2._connection, connection)
def test_uninstrument_connection(self):
connection = mock.Mock()
connection = MockConnectionEmpty()
# Set connection.database to avoid a failure because mock can't
# be concatenated
# pylint: disable=attribute-defined-outside-init
connection.database = "-"
connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
self.assertIs(connection2.__wrapped__, connection)
self.assertIs(connection2._connection, connection)
connection3 = dbapi.uninstrument_connection(connection2)
self.assertIs(connection3, connection)
@ -307,10 +309,12 @@ def mock_connect(*args, **kwargs):
server_host = kwargs.get("server_host")
server_port = kwargs.get("server_port")
user = kwargs.get("user")
return MockConnection(database, server_port, server_host, user)
return MockConnectionWithAttributes(
database, server_port, server_host, user
)
class MockConnection:
class MockConnectionWithAttributes:
def __init__(self, database, server_port, server_host, user):
self.database = database
self.server_port = server_port
@ -343,3 +347,7 @@ class MockCursor:
def callproc(self, query, params=None, throw_exception=False):
if throw_exception:
raise Exception("Test Exception")
class MockConnectionEmpty:
pass