From f67eabc09c6304397532aa587a19b0809c1f3c1d Mon Sep 17 00:00:00 2001 From: Paulo Vital Date: Wed, 30 Aug 2023 15:49:16 +0200 Subject: [PATCH] test: Add tests for CTX manager usage on PEP0249. Add the test_connect_ctx_mngr UT to test the span creation when using the psycopg2 connect with context manager, and the test_connect_cursor_ctx_mngr UT when using the psycopg2 connect and cursos context manager. Signed-off-by: Paulo Vital --- tests/clients/test_psycopg2.py | 51 ++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/tests/clients/test_psycopg2.py b/tests/clients/test_psycopg2.py index 052d957c7..0932cb1c0 100644 --- a/tests/clients/test_psycopg2.py +++ b/tests/clients/test_psycopg2.py @@ -248,3 +248,54 @@ def test_register_type(self): ext.register_type(ext.UUID, self.cursor) ext.register_type(ext.UUIDARRAY, self.cursor) + def test_connect_ctx_mngr(self): + with tracer.start_active_span('test'): + with self.db as conn: + cursor = conn.cursor() + cursor.execute("""SELECT * from users""") + cursor.fetchone() + + spans = self.recorder.queued_spans() + self.assertEqual(2, len(spans)) + + db_span = spans[0] + test_span = spans[1] + + self.assertEqual("test", test_span.data["sdk"]["name"]) + self.assertEqual(test_span.t, db_span.t) + self.assertEqual(db_span.p, test_span.s) + + self.assertEqual(None, db_span.ec) + + self.assertEqual(db_span.n, "postgres") + self.assertEqual(db_span.data["pg"]["db"], testenv['postgresql_db']) + self.assertEqual(db_span.data["pg"]["user"], testenv['postgresql_user']) + self.assertEqual(db_span.data["pg"]["stmt"], 'SELECT * from users') + self.assertEqual(db_span.data["pg"]["host"], testenv['postgresql_host']) + self.assertEqual(db_span.data["pg"]["port"], testenv['postgresql_port']) + + def test_connect_cursor_ctx_mngr(self): + with tracer.start_active_span('test'): + with self.db as conn: + with conn.cursor() as cursor: + cursor.execute("""SELECT * from users""") + cursor.fetchone() + + spans = self.recorder.queued_spans() + self.assertEqual(2, len(spans)) + + db_span = spans[0] + test_span = spans[1] + + self.assertEqual("test", test_span.data["sdk"]["name"]) + self.assertEqual(test_span.t, db_span.t) + self.assertEqual(db_span.p, test_span.s) + + self.assertEqual(None, db_span.ec) + + self.assertEqual(db_span.n, "postgres") + self.assertEqual(db_span.data["pg"]["db"], testenv['postgresql_db']) + self.assertEqual(db_span.data["pg"]["user"], testenv['postgresql_user']) + self.assertEqual(db_span.data["pg"]["stmt"], 'SELECT * from users') + self.assertEqual(db_span.data["pg"]["host"], testenv['postgresql_host']) + self.assertEqual(db_span.data["pg"]["port"], testenv['postgresql_port'])