fix(main) Increase code coverage to 94%

Changelog:
- Add unit tests for proxy
- Add test case to connector for drivers that require the local proxy
- Make proper adjustments to code

Author: Uziel Silva

google/cloud/sql/connector/connector.py

@@ -40,7 +40,7 @@
 from google.cloud.sql.connector.lazy import LazyRefreshCache
 from google.cloud.sql.connector.monitored_cache import MonitoredCache
 import google.cloud.sql.connector.pg8000 as pg8000
-from google.cloud.sql.connector.proxy import start_local_proxy
+import google.cloud.sql.connector.proxy as proxy
 import google.cloud.sql.connector.psycopg as psycopg
 import google.cloud.sql.connector.pymysql as pymysql
 import google.cloud.sql.connector.pytds as pytds
@@ -454,7 +454,7 @@ async def connect_async(
             if driver in LOCAL_PROXY_DRIVERS:
                 local_socket_path = kwargs.pop("local_socket_path", "/tmp/connector-socket")
                 host = local_socket_path
-                self._proxy = start_local_proxy(
+                self._proxy = proxy.start_local_proxy(
                     sock,
                     socket_path=f"{local_socket_path}/.s.PGSQL.{SERVER_PROXY_PORT}",
                     loop=self._loop

google/cloud/sql/connector/proxy.py

@@ -74,9 +74,9 @@ def start_local_proxy(
 async def local_communication(
   unix_socket, ssl_sock, socket_path, loop
 ):
+    client, _ = await loop.sock_accept(unix_socket)
+    
     try:
-        client, _ = await loop.sock_accept(unix_socket)
-
         while True:
             data = await loop.sock_recv(client, LOCAL_PROXY_MAX_MESSAGE_SIZE)
             if not data:
@@ -85,8 +85,6 @@ async def local_communication(
             ssl_sock.sendall(data)
             response = ssl_sock.recv(LOCAL_PROXY_MAX_MESSAGE_SIZE)
             await loop.sock_sendall(client, response)
-    except Exception:
-        pass
     finally:
         client.close()
         os.remove(socket_path) # Clean up the socket file

tests/unit/test_connector.py

@@ -17,6 +17,8 @@
 import asyncio
 import os
 from threading import Thread
+import socket
+import ssl
 from typing import Union
 
 from aiohttp import ClientResponseError
@@ -35,6 +37,7 @@
 from google.cloud.sql.connector.exceptions import IncompatibleDriverError
 from google.cloud.sql.connector.instance import RefreshAheadCache
 from google.cloud.sql.connector.resolver import DnsResolver
+from google.cloud.sql.connector.proxy import start_local_proxy
 
 
 @pytest.mark.asyncio
@@ -283,6 +286,42 @@ async def test_Connector_connect_async(
             # verify connector made connection call
             assert connection is True
 
+@pytest.mark.usefixtures("proxy_server")
+@pytest.mark.asyncio
+async def test_Connector_connect_local_proxy(
+    fake_credentials: Credentials, fake_client: CloudSQLClient, context: ssl.SSLContext
+) -> None:
+    """Test that Connector.connect can launch start_local_proxy."""
+    async with Connector(
+        credentials=fake_credentials, loop=asyncio.get_running_loop()
+    ) as connector:
+        connector._client = fake_client
+        socket_path = "/tmp/connector-socket/socket"
+        ip_addr = "127.0.0.1"
+        ssl_sock = context.wrap_socket(
+            socket.create_connection((ip_addr, 3307)),
+            server_hostname=ip_addr,
+        )
+        loop = asyncio.get_running_loop()
+        task = start_local_proxy(ssl_sock, socket_path, loop)
+        # patch db connection creation
+        with patch("google.cloud.sql.connector.proxy.start_local_proxy") as mock_proxy:
+            with patch("google.cloud.sql.connector.psycopg.connect") as mock_connect:
+                mock_connect.return_value = True
+                mock_proxy.return_value = task
+                connection = await connector.connect_async(
+                    "test-project:test-region:test-instance",
+                    "psycopg",
+                    user="my-user",
+                    password="my-pass",
+                    db="my-db",
+                    local_socket_path=socket_path,
+                )
+                # verify connector called local proxy
+                mock_connect.assert_called_once()
+                mock_proxy.assert_called_once()
+                assert connection is True
+
 
 @pytest.mark.asyncio
 async def test_Connector_connect_async_multiple_event_loops(

tests/unit/test_proxy.py

@@ -0,0 +1,75 @@
+"""
+Copyright 2025 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 asyncio
+import socket
+import ssl
+from typing import Any
+
+from mock import Mock
+import pytest
+
+from google.cloud.sql.connector import proxy
+
+LOCAL_PROXY_MAX_MESSAGE_SIZE = 10485760
+
+@pytest.mark.usefixtures("proxy_server")
+@pytest.mark.asyncio
+async def test_proxy_creates_folder(context: ssl.SSLContext, kwargs: Any) -> None:
+    """Test to verify that the proxy server is getting back the task."""
+    ip_addr = "127.0.0.1"
+    path = "/tmp/connector-socket/socket"
+    sock = context.wrap_socket(
+        socket.create_connection((ip_addr, 3307)),
+        server_hostname=ip_addr,
+    )
+    loop = asyncio.get_running_loop()
+
+    task = proxy.start_local_proxy(sock, path, loop)
+    assert (task is not None)
+    
+    proxy_task = asyncio.gather(task)
+    try:
+        await asyncio.wait_for(proxy_task, timeout=0.1)
+    except TimeoutError:
+        pass # This task runs forever so it is expected to throw this exception
+
+@pytest.mark.usefixtures("proxy_server")
+@pytest.mark.asyncio
+async def test_local_proxy_communication(context: ssl.SSLContext, kwargs: Any) -> None:
+    """Test to verify that the communication is getting through."""
+    socket_path = "/tmp/connector-socket/socket"
+    ssl_sock = Mock(spec=ssl.SSLSocket)
+    loop = asyncio.get_running_loop()
+
+    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
+        ssl_sock.recv.return_value = b"Received"
+
+        task = proxy.start_local_proxy(ssl_sock, socket_path, loop)
+        
+        client.connect(socket_path)
+        client.sendall(b"Test")
+        await asyncio.sleep(1)
+        
+        ssl_sock.sendall.assert_called_with(b"Test")
+        response = client.recv(LOCAL_PROXY_MAX_MESSAGE_SIZE)
+        assert (response == b"Received")
+
+        client.close()
+        await asyncio.sleep(1)
+
+        proxy_task = asyncio.gather(task)
+        await asyncio.wait_for(proxy_task, timeout=2)