From ccf7ce05a22c8d6896277e05d84973259594262d 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_mysqlclient.py | 90 +++++++++++++++++++++++++++++++ tests/clients/test_psycopg2.py | 77 ++++++++++++++++++++++++++ tests/clients/test_pymysql.py | 90 +++++++++++++++++++++++++++++++ 3 files changed, 257 insertions(+) diff --git a/tests/clients/test_mysqlclient.py b/tests/clients/test_mysqlclient.py index 02cf4c42f..3ac75eb52 100644 --- a/tests/clients/test_mysqlclient.py +++ b/tests/clients/test_mysqlclient.py @@ -217,3 +217,93 @@ def test_error_capture(self): self.assertEqual(db_span.data["mysql"]["stmt"], 'SELECT * from blah') self.assertEqual(db_span.data["mysql"]["host"], testenv['mysql_host']) self.assertEqual(db_span.data["mysql"]["port"], testenv['mysql_port']) + + def test_connect_ctx_mngr(self): + result = None + + with tracer.start_active_span('test'): + with self.db as conn: + cursor = conn.cursor() + result = cursor.execute("""SELECT * from users""") + cursor.fetchone() + + assert(result >= 0) + + 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, "mysql") + self.assertEqual(db_span.data["mysql"]["db"], testenv['mysql_db']) + self.assertEqual(db_span.data["mysql"]["user"], testenv['mysql_user']) + self.assertEqual(db_span.data["mysql"]["stmt"], 'SELECT * from users') + self.assertEqual(db_span.data["mysql"]["host"], testenv['mysql_host']) + self.assertEqual(db_span.data["mysql"]["port"], testenv['mysql_port']) + + def test_cursor_ctx_mngr(self): + result = None + + with tracer.start_active_span('test'): + with self.db.cursor() as cursor: + result = cursor.execute("""SELECT * from users""") + cursor.fetchone() + self.db.close() + + assert(result >= 0) + + 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, "mysql") + self.assertEqual(db_span.data["mysql"]["db"], testenv['mysql_db']) + self.assertEqual(db_span.data["mysql"]["user"], testenv['mysql_user']) + self.assertEqual(db_span.data["mysql"]["stmt"], 'SELECT * from users') + self.assertEqual(db_span.data["mysql"]["host"], testenv['mysql_host']) + self.assertEqual(db_span.data["mysql"]["port"], testenv['mysql_port']) + + def test_connect_cursor_ctx_mngr(self): + result = None + + with tracer.start_active_span('test'): + with self.db as conn: + with conn.cursor() as cursor: + result = cursor.execute("""SELECT * from users""") + cursor.fetchone() + + assert(result >= 0) + + 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, "mysql") + self.assertEqual(db_span.data["mysql"]["db"], testenv['mysql_db']) + self.assertEqual(db_span.data["mysql"]["user"], testenv['mysql_user']) + self.assertEqual(db_span.data["mysql"]["stmt"], 'SELECT * from users') + self.assertEqual(db_span.data["mysql"]["host"], testenv['mysql_host']) + self.assertEqual(db_span.data["mysql"]["port"], testenv['mysql_port']) diff --git a/tests/clients/test_psycopg2.py b/tests/clients/test_psycopg2.py index 052d957c7..a3f83e4e1 100644 --- a/tests/clients/test_psycopg2.py +++ b/tests/clients/test_psycopg2.py @@ -248,3 +248,80 @@ 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_cursor_ctx_mngr(self): + with tracer.start_active_span('test'): + with self.db.cursor() as cursor: + cursor.execute("""SELECT * from users""") + cursor.fetchone() + self.db.close() + + 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']) diff --git a/tests/clients/test_pymysql.py b/tests/clients/test_pymysql.py index bbc8106bb..4e333ba41 100644 --- a/tests/clients/test_pymysql.py +++ b/tests/clients/test_pymysql.py @@ -243,3 +243,93 @@ def test_error_capture(self): self.assertEqual(db_span.data["mysql"]["stmt"], 'SELECT * from blah') self.assertEqual(db_span.data["mysql"]["host"], testenv['mysql_host']) self.assertEqual(db_span.data["mysql"]["port"], testenv['mysql_port']) + + def test_connect_ctx_mngr(self): + result = None + + with tracer.start_active_span('test'): + with self.db as conn: + cursor = conn.cursor() + result = cursor.execute("""SELECT * from users""") + cursor.fetchone() + + assert(result >= 0) + + 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, "mysql") + self.assertEqual(db_span.data["mysql"]["db"], testenv['mysql_db']) + self.assertEqual(db_span.data["mysql"]["user"], testenv['mysql_user']) + self.assertEqual(db_span.data["mysql"]["stmt"], 'SELECT * from users') + self.assertEqual(db_span.data["mysql"]["host"], testenv['mysql_host']) + self.assertEqual(db_span.data["mysql"]["port"], testenv['mysql_port']) + + def test_cursor_ctx_mngr(self): + result = None + + with tracer.start_active_span('test'): + with self.db.cursor() as cursor: + result = cursor.execute("""SELECT * from users""") + cursor.fetchone() + self.db.close() + + assert(result >= 0) + + 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, "mysql") + self.assertEqual(db_span.data["mysql"]["db"], testenv['mysql_db']) + self.assertEqual(db_span.data["mysql"]["user"], testenv['mysql_user']) + self.assertEqual(db_span.data["mysql"]["stmt"], 'SELECT * from users') + self.assertEqual(db_span.data["mysql"]["host"], testenv['mysql_host']) + self.assertEqual(db_span.data["mysql"]["port"], testenv['mysql_port']) + + def test_connect_cursor_ctx_mngr(self): + result = None + + with tracer.start_active_span('test'): + with self.db as conn: + with conn.cursor() as cursor: + result = cursor.execute("""SELECT * from users""") + cursor.fetchone() + + assert(result >= 0) + + 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, "mysql") + self.assertEqual(db_span.data["mysql"]["db"], testenv['mysql_db']) + self.assertEqual(db_span.data["mysql"]["user"], testenv['mysql_user']) + self.assertEqual(db_span.data["mysql"]["stmt"], 'SELECT * from users') + self.assertEqual(db_span.data["mysql"]["host"], testenv['mysql_host']) + self.assertEqual(db_span.data["mysql"]["port"], testenv['mysql_port'])