From 2090ec7c916e08eeefdab22a3d5e33e59e772882 Mon Sep 17 00:00:00 2001
From: Pablo Martin <pablo.martin@lolamarket.com>
Date: Mon, 24 Apr 2023 13:41:15 +0200
Subject: [PATCH] Refactor setup and teardown of connection and test table into
 a fixture.

---
 tests/test_integration/test_utils.py | 86 +++++++++++-----------------
 1 file changed, 32 insertions(+), 54 deletions(-)

diff --git a/tests/test_integration/test_utils.py b/tests/test_integration/test_utils.py
index ba6023a..e195626 100644
--- a/tests/test_integration/test_utils.py
+++ b/tests/test_integration/test_utils.py
@@ -1,3 +1,5 @@
+import pytest
+
 from lolafect.lolaconfig import build_lolaconfig
 from lolafect.connections import (
     open_ssh_tunnel_with_s3_pkey,
@@ -25,7 +27,8 @@ from lolafect.utils import begin_sql_transaction, end_sql_transaction
 TEST_LOLACONFIG = build_lolaconfig(flow_name="testing-suite")
 
 
-def test_sql_transaction_persists_changes_properly():
+@pytest.fixture
+def connection_with_test_table():
     test_local_bind_host = "127.0.0.1"
     test_local_bind_port = 12345
 
@@ -44,26 +47,41 @@ def test_sql_transaction_persists_changes_properly():
             tunnel=tunnel
         ),
     )
-
     cursor = connection.cursor()
     cursor.execute("""
-    CREATE TABLE sandbox.lolafect_transaction_test_table
-    (
-        a_test_column INT
-    )
-    """)
+        CREATE TABLE sandbox.lolafect_transaction_test_table
+        (
+            a_test_column INT
+        )
+        """)
+
+    # Connection and table ready for tests
+    yield connection  # Test happens now
+    # Test finished, time to remove stuff and close connection
+
+    cursor.execute("""
+        DROP TABLE sandbox.lolafect_transaction_test_table
+        """
+                   )
+    close_mysql_connection.run(connection=connection)
+    close_ssh_tunnel.run(tunnel=tunnel)
+
+
+def test_sql_transaction_persists_changes_properly(connection_with_test_table):
+    cursor = connection_with_test_table.cursor()
+
     cursor.execute("""
     SELECT a_test_column
     FROM sandbox.lolafect_transaction_test_table
     """)
     table_is_empty_at_first = not bool(cursor.fetchall())  # An empty tuple yields False
 
-    begin_sql_transaction.run(connection=connection)
+    begin_sql_transaction.run(connection=connection_with_test_table)
     cursor.execute("""
     INSERT INTO sandbox.lolafect_transaction_test_table (a_test_column)
     VALUES (1)
     """)
-    end_sql_transaction.run(connection=connection, dry_run=False)
+    end_sql_transaction.run(connection=connection_with_test_table, dry_run=False)
 
     cursor.execute("""
     SELECT a_test_column
@@ -71,69 +89,29 @@ def test_sql_transaction_persists_changes_properly():
     """)
     table_has_a_record_after_transaction = bool(cursor.fetchall())  # A non-empty tuple yields True
 
-    cursor.execute("""
-    DROP TABLE sandbox.lolafect_transaction_test_table
-    """
-                   )
-
-    close_mysql_connection.run(connection=connection)
-    close_ssh_tunnel.run(tunnel=tunnel)
-
     assert table_is_empty_at_first and table_has_a_record_after_transaction
 
 
-def test_sql_transaction_rollbacks_changes_properly():
-    test_local_bind_host = "127.0.0.1"
-    test_local_bind_port = 12345
+def test_sql_transaction_rollbacks_changes_properly(connection_with_test_table):
+    cursor = connection_with_test_table.cursor()
 
-    tunnel = open_ssh_tunnel_with_s3_pkey.run(
-        s3_bucket_name=TEST_LOLACONFIG.S3_BUCKET_NAME,
-        ssh_tunnel_credentials=TEST_LOLACONFIG.SSH_TUNNEL_CREDENTIALS,
-        remote_target_host=TEST_LOLACONFIG.DW_CREDENTIALS["host"],
-        remote_target_port=TEST_LOLACONFIG.DW_CREDENTIALS["port"],
-        local_bind_host=test_local_bind_host,
-        local_bind_port=test_local_bind_port,
-    )
-
-    connection = connect_to_mysql.run(
-        mysql_credentials=TEST_LOLACONFIG.DW_CREDENTIALS,
-        overriding_host_and_port=get_local_bind_address_from_ssh_tunnel.run(
-            tunnel=tunnel
-        ),
-    )
-
-    cursor = connection.cursor()
-    cursor.execute("""
-    CREATE TABLE sandbox.lolafect_transaction_test_table
-    (
-        a_test_column INT
-    )
-    """)
     cursor.execute("""
     SELECT a_test_column
     FROM sandbox.lolafect_transaction_test_table
     """)
     table_is_empty_at_first = not bool(cursor.fetchall())  # An empty tuple yields False
 
-    begin_sql_transaction.run(connection=connection)
+    begin_sql_transaction.run(connection=connection_with_test_table)
     cursor.execute("""
     INSERT INTO sandbox.lolafect_transaction_test_table (a_test_column)
     VALUES (1)
     """)
-    end_sql_transaction.run(connection=connection, dry_run=True)
+    end_sql_transaction.run(connection=connection_with_test_table, dry_run=True)
 
     cursor.execute("""
     SELECT a_test_column
     FROM sandbox.lolafect_transaction_test_table
     """)
-    table_is_empty_after_rollback = not bool(cursor.fetchall())  # An tuple yields False
-
-    cursor.execute("""
-    DROP TABLE sandbox.lolafect_transaction_test_table
-    """
-                   )
-
-    close_mysql_connection.run(connection=connection)
-    close_ssh_tunnel.run(tunnel=tunnel)
+    table_is_empty_after_rollback = not bool(cursor.fetchall())  # A tuple yields False
 
     assert table_is_empty_at_first and table_is_empty_after_rollback