Close mysql connections and cursors in tests (#352)

Author: alexmojaki

tests/otel_integrations/test_mysql.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import sys
+
 import mysql.connector
 import pytest
 from dirty_equals import IsInt
@@ -10,6 +12,8 @@
 import logfire
 from logfire.testing import TestExporter
 
+pytestmark = pytest.mark.skipif(sys.version_info < (3, 9), reason='MySQL testcontainers has problems in 3.8')
+
 
 @pytest.fixture(scope='module')
 def mysql_container():
@@ -26,10 +30,12 @@ def get_mysql_connection(mysql_container: MySqlContainer):
 
 def test_mysql_instrumentation(exporter: TestExporter, mysql_container: MySqlContainer):
     logfire.instrument_mysql()
-    conn = get_mysql_connection(mysql_container)
-    cursor = conn.cursor()
-    cursor.execute('DROP TABLE IF EXISTS test')
-    cursor.execute('CREATE TABLE test (id INT PRIMARY KEY, name VARCHAR(255))')
+
+    with get_mysql_connection(mysql_container) as conn:
+        with conn.cursor() as cursor:
+            cursor.execute('DROP TABLE IF EXISTS test')
+            cursor.execute('CREATE TABLE test (id INT PRIMARY KEY, name VARCHAR(255))')
+
     assert exporter.exported_spans_as_dict() == snapshot(
         [
             {
@@ -72,38 +78,41 @@ def test_mysql_instrumentation(exporter: TestExporter, mysql_container: MySqlCon
 
 
 def test_instrument_mysql_connection(exporter: TestExporter, mysql_container: MySqlContainer):
-    conn = get_mysql_connection(mysql_container)
-    cursor = conn.cursor()
-    cursor.execute('DROP TABLE IF EXISTS test')
-    cursor.execute('CREATE TABLE test (id INT PRIMARY KEY, name VARCHAR(255))')
-    assert exporter.exported_spans_as_dict() == []
-
-    conn = logfire.instrument_mysql(conn)
-    cursor = conn.cursor()
-    cursor.execute('INSERT INTO test (id, name) VALUES (1, "test")')
-    assert exporter.exported_spans_as_dict() == snapshot(
-        [
-            {
-                'name': 'INSERT',
-                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
-                'parent': None,
-                'start_time': 1000000000,
-                'end_time': 2000000000,
-                'attributes': {
-                    'logfire.span_type': 'span',
-                    'logfire.msg': 'INSERT INTO test (id, name) VALUES (1, "test")',
-                    'db.system': 'mysql',
-                    'db.name': 'test',
-                    'db.statement': 'INSERT INTO test (id, name) VALUES (1, "test")',
-                    'db.user': 'test',
-                    'net.peer.name': 'localhost',
-                    'net.peer.port': IsInt(),
-                },
-            }
-        ]
-    )
+    with get_mysql_connection(mysql_container) as conn:
+        with conn.cursor() as cursor:
+            cursor.execute('DROP TABLE IF EXISTS test')
+            cursor.execute('CREATE TABLE test (id INT PRIMARY KEY, name VARCHAR(255))')
+
+        assert exporter.exported_spans_as_dict() == []
+
+        conn = logfire.instrument_mysql(conn)
+        with conn.cursor() as cursor:
+            cursor.execute('INSERT INTO test (id, name) VALUES (1, "test")')
+
+        assert exporter.exported_spans_as_dict() == snapshot(
+            [
+                {
+                    'name': 'INSERT',
+                    'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                    'parent': None,
+                    'start_time': 1000000000,
+                    'end_time': 2000000000,
+                    'attributes': {
+                        'logfire.span_type': 'span',
+                        'logfire.msg': 'INSERT INTO test (id, name) VALUES (1, "test")',
+                        'db.system': 'mysql',
+                        'db.name': 'test',
+                        'db.statement': 'INSERT INTO test (id, name) VALUES (1, "test")',
+                        'db.user': 'test',
+                        'net.peer.name': 'localhost',
+                        'net.peer.port': IsInt(),
+                    },
+                }
+            ]
+        )
+
+        conn = MySQLInstrumentor().uninstrument_connection(conn)  # type: ignore
+        with conn.cursor() as cursor:  # type: ignore
+            cursor.execute('INSERT INTO test (id, name) VALUES (2, "test-2")')  # type: ignore
 
-    conn = MySQLInstrumentor().uninstrument_connection(conn)  # type: ignore
-    cursor = conn.cursor()  # type: ignore
-    cursor.execute('INSERT INTO test (id, name) VALUES (2, "test-2")')  # type: ignore
-    assert len(exporter.exported_spans_as_dict()) == 1
+        assert len(exporter.exported_spans_as_dict()) == 1