mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-07-29 13:12:39 +08:00
Refactor TracedConnectionProxy (#1097)
* Refactor TracedConnectionProxy Fixes #1077 * Remove unecessary methods
This commit is contained in:
@ -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
|
||||
|
Reference in New Issue
Block a user