Skip to content
Open
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ class MockConnection:

def __init__(self, *args, **kwargs):
self.cursor_factory = kwargs.pop("cursor_factory", None)
self._initialized = True

def cursor(self):
if self.cursor_factory:
Expand All @@ -63,6 +64,17 @@ def cursor(self):
def get_dsn_parameters(self): # pylint: disable=no-self-use
return {"dbname": "test"}

def __setattr__(self, name, value):
"""
This method is overridden to prevent adding new attributes to the MockConnection object.
At runtime, the psycopg2 connection object does not have attribute
_is_instrumented_by_opentelemetry, it will throw AttributeError when setting this
attribute.
"""
if hasattr(self, '_initialized') and not hasattr(self, name):
raise AttributeError(f"Cannot add new attribute '{name}' to MockConnection")
super().__setattr__(name, value)


class TestPostgresqlIntegration(TestBase):
def setUp(self):
Expand Down