chore: update system tests to use Connector object (#392)

Author: jackwotherspoon

tests/system/test_ip_types.py

@@ -13,19 +13,20 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-import logging
 import os
 import uuid
 
 import pymysql
 import pytest
 import sqlalchemy
-from google.cloud.sql.connector import connector, IPTypes
+from google.cloud.sql.connector import Connector, IPTypes
 
 table_name = f"books_{uuid.uuid4().hex}"
 
 
-def init_connection_engine(ip_type: IPTypes) -> sqlalchemy.engine.Engine:
+def init_connection_engine(
+    connector: Connector, ip_type: IPTypes
+) -> sqlalchemy.engine.Engine:
     def getconn() -> pymysql.connections.Connection:
         conn: pymysql.connections.Connection = connector.connect(
             os.environ["MYSQL_CONNECTION_NAME"],
@@ -45,19 +46,15 @@ def getconn() -> pymysql.connections.Connection:
 
 
 def test_public_ip() -> None:
-    try:
-        pool = init_connection_engine(IPTypes.PUBLIC)
-    except Exception as e:
-        logging.exception("Failed to initialize pool with public IP", e)
-    with pool.connect() as conn:
-        conn.execute("SELECT 1")
+    with Connector() as connector:
+        pool = init_connection_engine(connector, IPTypes.PUBLIC)
+        with pool.connect() as conn:
+            conn.execute("SELECT 1")
 
 
 @pytest.mark.private_ip
 def test_private_ip() -> None:
-    try:
-        pool = init_connection_engine(IPTypes.PRIVATE)
-    except Exception as e:
-        logging.exception("Failed to initialize pool with private IP", e)
-    with pool.connect() as conn:
-        conn.execute("SELECT 1")
+    with Connector() as connector:
+        pool = init_connection_engine(connector, IPTypes.PRIVATE)
+        with pool.connect() as conn:
+            conn.execute("SELECT 1")

tests/system/test_pg8000_connection.py

@@ -20,7 +20,7 @@
 import pg8000
 import pytest
 import sqlalchemy
-from google.cloud.sql.connector import connector
+from google.cloud.sql.connector import Connector
 
 table_name = f"books_{uuid.uuid4().hex}"
 
@@ -30,14 +30,15 @@
 # 'creator' argument to 'create_engine'
 def init_connection_engine() -> sqlalchemy.engine.Engine:
     def getconn() -> pg8000.dbapi.Connection:
-        conn: pg8000.dbapi.Connection = connector.connect(
-            os.environ["POSTGRES_CONNECTION_NAME"],
-            "pg8000",
-            user=os.environ["POSTGRES_USER"],
-            password=os.environ["POSTGRES_PASS"],
-            db=os.environ["POSTGRES_DB"],
-        )
-        return conn
+        with Connector() as connector:
+            conn: pg8000.dbapi.Connection = connector.connect(
+                os.environ["POSTGRES_CONNECTION_NAME"],
+                "pg8000",
+                user=os.environ["POSTGRES_USER"],
+                password=os.environ["POSTGRES_PASS"],
+                db=os.environ["POSTGRES_DB"],
+            )
+            return conn
 
     engine = sqlalchemy.create_engine(
         "postgresql+pg8000://",

tests/system/test_pg8000_iam_auth.py

@@ -20,7 +20,7 @@
 import pg8000
 import pytest
 import sqlalchemy
-from google.cloud.sql.connector import connector
+from google.cloud.sql.connector import Connector
 
 table_name = f"books_{uuid.uuid4().hex}"
 
@@ -30,14 +30,15 @@
 # 'creator' argument to 'create_engine'
 def init_connection_engine() -> sqlalchemy.engine.Engine:
     def getconn() -> pg8000.dbapi.Connection:
-        conn: pg8000.dbapi.Connection = connector.connect(
-            os.environ["POSTGRES_IAM_CONNECTION_NAME"],
-            "pg8000",
-            user=os.environ["POSTGRES_IAM_USER"],
-            db=os.environ["POSTGRES_DB"],
-            enable_iam_auth=True,
-        )
-        return conn
+        with Connector() as connector:
+            conn: pg8000.dbapi.Connection = connector.connect(
+                os.environ["POSTGRES_IAM_CONNECTION_NAME"],
+                "pg8000",
+                user=os.environ["POSTGRES_IAM_USER"],
+                db=os.environ["POSTGRES_DB"],
+                enable_iam_auth=True,
+            )
+            return conn
 
     engine = sqlalchemy.create_engine(
         "postgresql+pg8000://",

tests/system/test_pymysql_connection.py

@@ -20,7 +20,7 @@
 import pymysql
 import pytest
 import sqlalchemy
-from google.cloud.sql.connector import connector
+from google.cloud.sql.connector import Connector
 
 table_name = f"books_{uuid.uuid4().hex}"
 
@@ -30,14 +30,15 @@
 # 'creator' argument to 'create_engine'
 def init_connection_engine() -> sqlalchemy.engine.Engine:
     def getconn() -> pymysql.connections.Connection:
-        conn: pymysql.connections.Connection = connector.connect(
-            os.environ["MYSQL_CONNECTION_NAME"],
-            "pymysql",
-            user=os.environ["MYSQL_USER"],
-            password=os.environ["MYSQL_PASS"],
-            db=os.environ["MYSQL_DB"],
-        )
-        return conn
+        with Connector() as connector:
+            conn: pymysql.connections.Connection = connector.connect(
+                os.environ["MYSQL_CONNECTION_NAME"],
+                "pymysql",
+                user=os.environ["MYSQL_USER"],
+                password=os.environ["MYSQL_PASS"],
+                db=os.environ["MYSQL_DB"],
+            )
+            return conn
 
     engine = sqlalchemy.create_engine(
         "mysql+pymysql://",

tests/system/test_pytds_connection.py

@@ -20,7 +20,7 @@
 import pytds
 import pytest
 import sqlalchemy
-from google.cloud.sql.connector import connector
+from google.cloud.sql.connector import Connector
 
 table_name = f"books_{uuid.uuid4().hex}"
 
@@ -31,14 +31,15 @@
 # 'sqlalchemy-pytds` in your dependencies
 def init_connection_engine() -> sqlalchemy.engine.Engine:
     def getconn() -> pytds.Connection:
-        conn = connector.connect(
-            os.environ["SQLSERVER_CONNECTION_NAME"],
-            "pytds",
-            user=os.environ["SQLSERVER_USER"],
-            password=os.environ["SQLSERVER_PASS"],
-            db=os.environ["SQLSERVER_DB"],
-        )
-        return conn
+        with Connector() as connector:
+            conn = connector.connect(
+                os.environ["SQLSERVER_CONNECTION_NAME"],
+                "pytds",
+                user=os.environ["SQLSERVER_USER"],
+                password=os.environ["SQLSERVER_PASS"],
+                db=os.environ["SQLSERVER_DB"],
+            )
+            return conn
 
     engine = sqlalchemy.create_engine(
         "mssql+pytds://localhost",