chore: improve tests with mocks and coverage (#238)

* chore: update refresh_utils tests to use mocks

* chore: add basic tests for Connector

Author: jackwotherspoon

noxfile.py

@@ -39,6 +39,7 @@ def lint(session):
         "sqlalchemy-stubs==0.4",
         "types-pkg-resources==0.1.3",
         "types-PyMySQL==1.0.6",
+        "types-mock==4.0.5",
         "twine==3.7.1",
     )
     session.install("-r", "requirements.txt")

tests/system/test_connector_object.py

@@ -0,0 +1,82 @@
+""""
+Copyright 2021 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import os
+import pymysql
+import sqlalchemy
+import logging
+import google.auth
+from google.cloud.sql.connector import connector
+
+
+def init_connection_engine(
+    custom_connector: connector.Connector,
+) -> sqlalchemy.engine.Engine:
+    def getconn() -> pymysql.connections.Connection:
+        conn = custom_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://",
+        creator=getconn,
+    )
+    return engine
+
+
+def test_connector_with_credentials() -> None:
+    """Test Connector object connection with credentials loaded from file."""
+    credentials, project = google.auth.load_credentials_from_file(
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"]
+    )
+    custom_connector = connector.Connector(credentials=credentials)
+    try:
+        pool = init_connection_engine(custom_connector)
+
+        with pool.connect() as conn:
+            conn.execute("SELECT 1")
+
+    except Exception as e:
+        logging.exception("Failed to connect with credentials from file!", e)
+
+
+def test_multiple_connectors() -> None:
+    """Test that same Cloud SQL instance can connect with two Connector objects."""
+    first_connector = connector.Connector()
+    second_connector = connector.Connector()
+    try:
+        pool = init_connection_engine(first_connector)
+        pool2 = init_connection_engine(second_connector)
+
+        with pool.connect() as conn:
+            conn.execute("SELECT 1")
+
+        with pool2.connect() as conn:
+            conn.execute("SELECT 1")
+
+        instance_connection_string = os.environ["MYSQL_CONNECTION_NAME"]
+        assert instance_connection_string in first_connector._instances
+        assert instance_connection_string in second_connector._instances
+        assert (
+            first_connector._instances[instance_connection_string]
+            != second_connector._instances[instance_connection_string]
+        )
+    except Exception as e:
+        logging.exception("Failed to connect with multiple Connector objects!", e)

tests/unit/test_connector.py

@@ -16,6 +16,7 @@
 
 import pytest  # noqa F401 Needed to run the tests
 from google.cloud.sql.connector.instance_connection_manager import (
+    IPTypes,
     InstanceConnectionManager,
 )
 from google.cloud.sql.connector import connector
@@ -53,3 +54,12 @@ async def timeout_stub(*args: Any, **kwargs: Any) -> None:
             "pymysql",
             timeout=timeout,
         )
+
+
+def test_default_Connector_Init() -> None:
+    """Test that default Connector __init__ sets properties properly."""
+    default_connector = connector.Connector()
+    assert default_connector._ip_type == IPTypes.PUBLIC
+    assert default_connector._enable_iam_auth is False
+    assert default_connector._timeout == 30
+    assert default_connector._credentials is None

tests/unit/test_instance_connection_manager.py

@@ -29,17 +29,18 @@
 
 
 @pytest.fixture
-def mock_credentials() -> Credentials:
+def credentials() -> Credentials:
     return Mock(spec=Credentials)
 
 
 @pytest.fixture
 def icm(
-    async_loop: asyncio.AbstractEventLoop, connect_string: str
+    async_loop: asyncio.AbstractEventLoop,
 ) -> InstanceConnectionManager:
     keys = asyncio.run_coroutine_threadsafe(generate_keys(), async_loop)
-    icm = InstanceConnectionManager(connect_string, "pymysql", keys, async_loop)
-
+    icm = InstanceConnectionManager(
+        "my-project:my-region:my-instance", "pymysql", keys, async_loop
+    )
     return icm
 
 
@@ -155,6 +156,12 @@ async def test_perform_refresh_replaces_invalid_result(
     # allow more frequent refreshes for tests
     setattr(icm, "_refresh_rate_limiter", test_rate_limiter)
 
+    # set current to valid MockMetadata instance
+    setattr(icm, "_get_instance_data", _get_metadata_success)
+    icm._current = asyncio.run_coroutine_threadsafe(
+        icm._perform_refresh(), icm._loop
+    ).result(timeout=10)
+
     # stub _get_instance_data to throw an error
     setattr(icm, "_get_instance_data", _get_metadata_error)
     icm._current = asyncio.run_coroutine_threadsafe(
@@ -197,7 +204,7 @@ async def test_force_refresh_cancels_pending_refresh(
 
 
 def test_auth_init_with_credentials_object(
-    icm: InstanceConnectionManager, mock_credentials: Credentials
+    icm: InstanceConnectionManager, credentials: Credentials
 ) -> None:
     """
     Test that InstanceConnectionManager's _auth_init initializes _credentials
@@ -207,22 +214,22 @@ def test_auth_init_with_credentials_object(
     with patch(
         "google.cloud.sql.connector.instance_connection_manager.with_scopes_if_required"
     ) as mock_auth:
-        mock_auth.return_value = mock_credentials
-        icm._auth_init(credentials=mock_credentials)
+        mock_auth.return_value = credentials
+        icm._auth_init(credentials=credentials)
         assert isinstance(icm._credentials, Credentials)
         mock_auth.assert_called_once()
 
 
 def test_auth_init_with_default_credentials(
-    icm: InstanceConnectionManager, mock_credentials: Credentials
+    icm: InstanceConnectionManager, credentials: Credentials
 ) -> None:
     """
     Test that InstanceConnectionManager's _auth_init initializes _credentials
     with application default credentials when credentials are not specified.
     """
     setattr(icm, "_credentials", None)
     with patch("google.auth.default") as mock_auth:
-        mock_auth.return_value = mock_credentials, None
+        mock_auth.return_value = credentials, None
         icm._auth_init(credentials=None)
         assert isinstance(icm._credentials, Credentials)
         mock_auth.assert_called_once()

tests/unit/test_refresh_utils.py

@@ -16,29 +16,75 @@
 from typing import Any
 
 import aiohttp
-import google.auth
+from google.auth.credentials import Credentials
+import json
 import pytest  # noqa F401 Needed to run the tests
+from mock import AsyncMock, Mock, patch
 
 from google.cloud.sql.connector.refresh_utils import _get_ephemeral, _get_metadata
 from google.cloud.sql.connector.utils import generate_keys
 
 
+class FakeClientSessionGet:
+    """Helper class to return mock data for get request."""
+
+    async def text(self) -> str:
+        response = {
+            "kind": "sql#connectSettings",
+            "serverCaCert": {
+                "kind": "sql#sslCert",
+                "certSerialNumber": "0",
+                "cert": "-----BEGIN CERTIFICATE-----\nabc123\n-----END CERTIFICATE-----",
+                "commonName": "Google",
+                "sha1Fingerprint": "abc",
+                "instance": "my-instance",
+                "createTime": "2021-10-18T18:48:03.785Z",
+                "expirationTime": "2031-10-16T18:49:03.785Z",
+            },
+            "ipAddresses": [
+                {"type": "PRIMARY", "ipAddress": "0.0.0.0"},
+                {"type": "PRIVATE", "ipAddress": "1.0.0.0"},
+            ],
+            "region": "my-region",
+            "databaseVersion": "MYSQL_8_0",
+            "backendType": "SECOND_GEN",
+        }
+        return json.dumps(response)
+
+
+class FakeClientSessionPost:
+    """Helper class to return mock data for post request."""
+
+    async def text(self) -> str:
+        response = {
+            "ephemeralCert": {
+                "kind": "sql#sslCert",
+                "certSerialNumber": "",
+                "cert": "-----BEGIN CERTIFICATE-----\nabc123\n-----END CERTIFICATE-----",
+            }
+        }
+        return json.dumps(response)
+
+
+@pytest.fixture
+def credentials() -> Credentials:
+    credentials = Mock(spec=Credentials)
+    credentials.valid = True
+    credentials.token = "12345"
+    return credentials
+
+
 @pytest.mark.asyncio
-async def test_get_ephemeral(connect_string: str) -> None:
+@patch("aiohttp.ClientSession.post", new_callable=AsyncMock)
+async def test_get_ephemeral(mock_post: AsyncMock, credentials: Credentials) -> None:
     """
-    Test to check whether _get_ephemeral runs without problems given a valid
-    connection string.
+    Test to check whether _get_ephemeral runs without problems given valid
+    parameters.
     """
+    mock_post.return_value = FakeClientSessionPost()
 
-    project = connect_string.split(":")[0]
-    instance = connect_string.split(":")[2]
-
-    credentials, project = google.auth.default(
-        scopes=[
-            "https://www.googleapis.com/auth/sqlservice.admin",
-            "https://www.googleapis.com/auth/cloud-platform",
-        ]
-    )
+    project = "my-project"
+    instance = "my-instance"
 
     _, pub_key = await generate_keys()
 
@@ -56,21 +102,16 @@ async def test_get_ephemeral(connect_string: str) -> None:
 
 
 @pytest.mark.asyncio
-async def test_get_metadata(connect_string: str) -> None:
+@patch("aiohttp.ClientSession.get", new_callable=AsyncMock)
+async def test_get_metadata(mock_get: AsyncMock, credentials: Credentials) -> None:
     """
-    Test to check whether _get_metadata runs without problems given a valid
-    connection string.
+    Test to check whether _get_metadata runs without problems given valid
+    parameters.
     """
+    mock_get.return_value = FakeClientSessionGet()
 
-    project = connect_string.split(":")[0]
-    instance = connect_string.split(":")[2]
-
-    credentials, project = google.auth.default(
-        scopes=[
-            "https://www.googleapis.com/auth/sqlservice.admin",
-            "https://www.googleapis.com/auth/cloud-platform",
-        ]
-    )
+    project = "my-project"
+    instance = "my-instance"
 
     async with aiohttp.ClientSession() as client_session:
         result = await _get_metadata(client_session, credentials, project, instance)