Add context manager tests to db modules that conform to PEP-249

Signed-off-by: Varsha GS <varsha.gs@ibm.com>

remove result

Signed-off-by: Varsha GS <varsha.gs@ibm.com>

tests/clients/test_mysqlclient.py

@@ -217,3 +217,80 @@ 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_cursor_ctx_mgr(self):
+        with tracer.start_active_span("test"):
+            with self.db as connection:
+                with connection.cursor() as cursor:
+                    cursor.execute("""SELECT * from users""")
+
+        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.assertIsNone(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_ctx_mgr(self):
+        with tracer.start_active_span("test"):
+            with self.db as connection:
+                cursor = connection.cursor()
+                cursor.execute("""SELECT * from users""")
+
+
+        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.assertIsNone(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_mgr(self):
+        with tracer.start_active_span("test"):
+            connection = self.db
+            with connection.cursor() as cursor:
+                cursor.execute("""SELECT * from users""")
+
+
+        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.assertIsNone(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"])

tests/clients/test_psycopg2.py

@@ -248,3 +248,77 @@ def test_register_type(self):
         ext.register_type(ext.UUID, self.cursor)
         ext.register_type(ext.UUIDARRAY, self.cursor)
 
+    def test_connect_cursor_ctx_mgr(self):
+        with tracer.start_active_span("test"):
+            with self.db as connection:
+                with connection.cursor() as cursor:
+                    cursor.execute("""SELECT * from users""")
+
+        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.assertIsNone(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_ctx_mgr(self):
+        with tracer.start_active_span("test"):
+            with self.db as connection:
+                cursor = connection.cursor()
+                cursor.execute("""SELECT * from users""")
+
+        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.assertIsNone(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_mgr(self):
+        with tracer.start_active_span("test"):
+            connection = self.db
+            with connection.cursor() as cursor:
+                cursor.execute("""SELECT * from users""")
+
+        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.assertIsNone(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"])

tests/clients/test_pymysql.py

@@ -243,3 +243,79 @@ 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_cursor_ctx_mgr(self):
+        with tracer.start_active_span("test"):
+            with self.db as connection:
+                with connection.cursor() as cursor:
+                    cursor.execute("""SELECT * from users""")
+
+        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.assertIsNone(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_ctx_mgr(self):
+        with tracer.start_active_span("test"):
+            with self.db as connection:
+                cursor = connection.cursor()
+                cursor.execute("""SELECT * from users""")
+
+        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.assertIsNone(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_mgr(self):
+        with tracer.start_active_span("test"):
+            connection = self.db
+            with connection.cursor() as cursor:
+                cursor.execute("""SELECT * from users""")
+
+
+        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.assertIsNone(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"])