test: Fix pymysql tests to make database content predicatable

Signed-off-by: Ferenc Géczi <[email protected]>

Author: Ferenc-

tests/clients/test_pymysql.py

@@ -12,56 +12,44 @@
 
 logger = logging.getLogger(__name__)
 
-create_table_query = 'CREATE TABLE IF NOT EXISTS users(id serial primary key, \
-                      name varchar(40) NOT NULL, email varchar(40) NOT NULL)'
-
-create_proc_query = """
-CREATE PROCEDURE test_proc(IN t VARCHAR(255))
-BEGIN
-    SELECT name FROM users WHERE name = t;
-END
-"""
-
-db = pymysql.connect(host=testenv['mysql_host'], port=testenv['mysql_port'],
-                     user=testenv['mysql_user'], passwd=testenv['mysql_pw'],
-                     db=testenv['mysql_db'])
-
-cursor = db.cursor()
-cursor.execute(create_table_query)
-
-while cursor.nextset() is not None:
-    pass
-
-cursor.execute('DROP PROCEDURE IF EXISTS test_proc')
-
-while cursor.nextset() is not None:
-    pass
-
-cursor.execute(create_proc_query)
-
-while cursor.nextset() is not None:
-    pass
-
-cursor.close()
-db.close()
-
 
 class TestPyMySQL(unittest.TestCase):
     def setUp(self):
         self.db = pymysql.connect(host=testenv['mysql_host'], port=testenv['mysql_port'],
                                   user=testenv['mysql_user'], passwd=testenv['mysql_pw'],
                                   db=testenv['mysql_db'])
+        database_setup_query = """
+        DROP TABLE IF EXISTS users; |
+        CREATE TABLE users(
+            id serial primary key,
+            name varchar(40) NOT NULL,
+            email varchar(40) NOT NULL
+        ); |
+        INSERT INTO users(name, email) VALUES('kermit', '[email protected]'); |
+        DROP PROCEDURE IF EXISTS test_proc; |
+        CREATE PROCEDURE test_proc(IN t VARCHAR(255))
+        BEGIN
+            SELECT name FROM users WHERE name = t;
+        END
+        """
+        setup_cursor = self.db.cursor()
+        for s in database_setup_query.split('|'):
+          setup_cursor.execute(s)
+
         self.cursor = self.db.cursor()
         self.recorder = tracer.recorder
         self.recorder.clear_spans()
         tracer.cur_ctx = None
 
     def tearDown(self):
-        """ Do nothing for now """
-        return None
+        if self.cursor and self.cursor.connection.open:
+          self.cursor.close()
+        if self.db and self.db.open:
+          self.db.close()
 
     def test_vanilla_query(self):
-        self.cursor.execute("""SELECT * from users""")
+        affected_rows = self.cursor.execute("""SELECT * from users""")
+        self.assertEqual(1, affected_rows)
         result = self.cursor.fetchone()
         self.assertEqual(3, len(result))
 
@@ -70,10 +58,11 @@ def test_vanilla_query(self):
 
     def test_basic_query(self):
         with tracer.start_active_span('test'):
-            result = self.cursor.execute("""SELECT * from users""")
-            self.cursor.fetchone()
+            affected_rows = self.cursor.execute("""SELECT * from users""")
+            result = self.cursor.fetchone()
 
-        self.assertTrue(result >= 0)
+        self.assertEqual(1, affected_rows)
+        self.assertEqual(3, len(result))
 
         spans = self.recorder.queued_spans()
         self.assertEqual(2, len(spans))
@@ -95,10 +84,11 @@ def test_basic_query(self):
 
     def test_query_with_params(self):
         with tracer.start_active_span('test'):
-            result = self.cursor.execute("""SELECT * from users where id=1""")
-            self.cursor.fetchone()
+            affected_rows = self.cursor.execute("""SELECT * from users where id=1""")
+            result = self.cursor.fetchone()
 
-        self.assertTrue(result >= 0)
+        self.assertEqual(1, affected_rows)
+        self.assertEqual(3, len(result))
 
         spans = self.recorder.queued_spans()
         self.assertEqual(2, len(spans))
@@ -120,11 +110,11 @@ def test_query_with_params(self):
 
     def test_basic_insert(self):
         with tracer.start_active_span('test'):
-            result = self.cursor.execute(
+            affected_rows = self.cursor.execute(
                         """INSERT INTO users(name, email) VALUES(%s, %s)""",
                         ('beaker', '[email protected]'))
 
-        self.assertEqual(1, result)
+        self.assertEqual(1, affected_rows)
 
         spans = self.recorder.queued_spans()
         self.assertEqual(2, len(spans))
@@ -146,11 +136,11 @@ def test_basic_insert(self):
 
     def test_executemany(self):
         with tracer.start_active_span('test'):
-            result = self.cursor.executemany("INSERT INTO users(name, email) VALUES(%s, %s)",
+            affected_rows  = self.cursor.executemany("INSERT INTO users(name, email) VALUES(%s, %s)",
                                              [('beaker', '[email protected]'), ('beaker', '[email protected]')])
             self.db.commit()
 
-        self.assertEqual(2, result)
+        self.assertEqual(2, affected_rows)
 
         spans = self.recorder.queued_spans()
         self.assertEqual(2, len(spans))
@@ -172,9 +162,9 @@ def test_executemany(self):
 
     def test_call_proc(self):
         with tracer.start_active_span('test'):
-            result = self.cursor.callproc('test_proc', ('beaker',))
+            callproc_result = self.cursor.callproc('test_proc', ('beaker',))
 
-        self.assertTrue(result)
+        self.assertIsInstance(callproc_result, tuple)
 
         spans = self.recorder.queued_spans()
         self.assertEqual(2, len(spans))
@@ -195,15 +185,14 @@ def test_call_proc(self):
         self.assertEqual(db_span.data["mysql"]["port"], testenv['mysql_port'])
 
     def test_error_capture(self):
-        result = None
+        affected_rows = None
         try:
             with tracer.start_active_span('test'):
-                result = self.cursor.execute("""SELECT * from blah""")
-                self.cursor.fetchone()
+                affected_rows = self.cursor.execute("""SELECT * from blah""")
         except Exception:
             pass
 
-        self.assertIsNone(result)
+        self.assertIsNone(affected_rows)
 
         spans = self.recorder.queued_spans()
         self.assertEqual(2, len(spans))
@@ -228,8 +217,9 @@ 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""")
+                    affected_rows = cursor.execute("""SELECT * from users""")
 
+        self.assertEqual(1, affected_rows)
         spans = self.recorder.queued_spans()
         self.assertEqual(2, len(spans))