1212# See the License for the specific language governing permissions and
1313# limitations under the License.
1414
15+ import psycopg
1516import psycopg2
1617from test_psycopg_functional import (
1718 POSTGRES_DB_NAME ,
2122 POSTGRES_USER ,
2223)
2324
25+ from opentelemetry .instrumentation .psycopg import PsycopgInstrumentor
2426from opentelemetry .instrumentation .psycopg2 import Psycopg2Instrumentor
2527from opentelemetry .test .test_base import TestBase
2628
2729
28- class TestFunctionalPsycopg (TestBase ):
30+ class TestFunctionalPsycopg2 (TestBase ):
2931 def setUp (self ):
3032 super ().setUp ()
3133 self ._tracer = self .tracer_provider .get_tracer (__name__ )
@@ -52,3 +54,49 @@ def test_commenter_enabled(self):
5254 self ._cursor .query .decode ("ascii" ),
5355 r"SELECT 1 /\*db_driver='psycopg2(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),libpq_version=\d*,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;" ,
5456 )
57+
58+
59+ class TestFunctionalPsycopg (TestBase ):
60+ def setUp (self ):
61+ super ().setUp ()
62+ PsycopgInstrumentor ().instrument (enable_commenter = True )
63+
64+ def tearDown (self ):
65+ PsycopgInstrumentor ().uninstrument ()
66+ super ().tearDown ()
67+
68+ def test_commenter_enabled_root_module (self ):
69+ connection = psycopg .connect (
70+ dbname = POSTGRES_DB_NAME ,
71+ user = POSTGRES_USER ,
72+ password = POSTGRES_PASSWORD ,
73+ host = POSTGRES_HOST ,
74+ port = POSTGRES_PORT ,
75+ )
76+ cursor = connection .cursor ()
77+ cursor .execute ("SELECT 1;" )
78+ self .assertRegex (
79+ cursor ._query .query .decode ("ascii" ),
80+ r"SELECT 1 /\*db_driver='psycopg(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),libpq_version=\d*,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;" ,
81+ )
82+
83+ cursor .close ()
84+ connection .close ()
85+
86+ def test_commenter_enabled_not_root_module (self ):
87+ connection = psycopg .Connection .connect (
88+ dbname = POSTGRES_DB_NAME ,
89+ user = POSTGRES_USER ,
90+ password = POSTGRES_PASSWORD ,
91+ host = POSTGRES_HOST ,
92+ port = POSTGRES_PORT ,
93+ )
94+ cursor = connection .cursor ()
95+ cursor .execute ("SELECT 1;" )
96+ self .assertRegex (
97+ cursor ._query .query .decode ("ascii" ),
98+ r"SELECT 1 /\*db_driver='psycopg(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;" ,
99+ )
100+
101+ cursor .close ()
102+ connection .close ()
0 commit comments