fix: increase test coverage

Author: Thameez Bodhanya

tests/unit/test_connector.py

@@ -17,6 +17,8 @@
 import socket
 from threading import Thread
 from typing import Union
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
 
 from aiohttp import ClientResponseError
 from google.auth.credentials import Credentials
@@ -28,6 +30,7 @@
 from google.cloud.sql.connector import IPTypes
 from google.cloud.sql.connector.client import CloudSQLClient
 from google.cloud.sql.connector.connection_name import ConnectionName
+from google.cloud.sql.connector.connector import ConnectorSocketFactory
 from google.cloud.sql.connector.exceptions import ClosedConnectorError
 from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError
 from google.cloud.sql.connector.exceptions import ConnectorLoopError
@@ -550,6 +553,57 @@ async def test_Connector_start_unix_socket_proxy_async(
         )
 
 
+@pytest.mark.asyncio
+async def test_Connector_start_unix_socket_proxy_async_rejects_duplicate_socket_path(
+    fake_credentials: Credentials,
+    ) -> None:
+    socket_path = "/tmp/cloudsql-test.sock"
+    async with Connector(
+        credentials=fake_credentials, loop=asyncio.get_running_loop()
+    ) as connector:
+        existing_proxy = MagicMock()
+        existing_proxy.unix_socket_path = socket_path
+        existing_proxy.close = AsyncMock()
+        connector._proxies.append(existing_proxy)
+
+        with pytest.raises(ValueError) as exc_info:
+            await connector.start_unix_socket_proxy_async(
+                "test-project:test-region:test-instance",
+                socket_path,
+            )
+
+        assert (
+            exc_info.value.args[0]
+            == f"Proxy for socket path {socket_path} already exists."
+        )
+
+
+@pytest.mark.asyncio
+async def test_Connector_close_async_closes_proxies_client_and_cache(
+    fake_credentials: Credentials,
+) -> None:
+    async with Connector(
+        credentials=fake_credentials, loop=asyncio.get_running_loop()
+    ) as connector:
+        proxy_instance = MagicMock()
+        proxy_instance.close = AsyncMock()
+        connector._proxies.append(proxy_instance)
+
+        connector._client = MagicMock()
+        connector._client.close = AsyncMock()
+
+        cached = MagicMock()
+        cached.close = AsyncMock()
+        connector._cache[("test-project:test-region:test-instance", False)] = cached
+
+        await connector.close_async()
+
+        assert connector._closed is True
+        proxy_instance.close.assert_awaited_once()
+        connector._client.close.assert_awaited_once()
+        cached.close.assert_awaited_once()
+
+
 def test_connect_closed_connector(
     fake_credentials: Credentials, fake_client: CloudSQLClient
 ) -> None:
@@ -678,7 +732,53 @@ async def test_Connector_connect_async_custom_dns_resolver_fallback(
                     # Restore original IPs
                     fake_client.instance.ip_addrs = original_ips
 
-class TestProtocol(asyncio.Protocol):
+
+@pytest.mark.asyncio
+async def test_Connector_get_cache_invalidates_bad_cached_entry(
+    fake_credentials: Credentials,
+    ) -> None:
+    connect_string = "test-project:test-region:test-instance"
+    async with Connector(
+        credentials=fake_credentials, loop=asyncio.get_running_loop()
+    ) as connector:
+        monitored_cache = MagicMock()
+        monitored_cache.closed = False
+        monitored_cache.close = AsyncMock()
+
+        conn_info = MagicMock()
+        conn_info.get_preferred_ip.side_effect = RuntimeError("invalid ip")
+        monitored_cache.connect_info = AsyncMock(return_value=conn_info)
+        connector._cache[(connect_string, False)] = monitored_cache
+
+        with (
+            patch.object(
+                connector._resolver,
+                "resolve",
+                AsyncMock(
+                    return_value=ConnectionName(
+                        "test-project",
+                        "test-region",
+                        "test-instance",
+                    )
+                ),
+            ),
+            patch.object(
+                connector, "_remove_cached", AsyncMock()
+            ) as mock_remove_cached,
+        ):
+            with pytest.raises(RuntimeError) as exc_info:
+                await connector._get_cache(
+                    connect_string,
+                    False,
+                    IPTypes.PUBLIC,
+                    None,
+                )
+
+        assert exc_info.value.args[0] == "invalid ip"
+        mock_remove_cached.assert_awaited_once_with(connect_string, False)
+
+
+class SocketTestProtocol(asyncio.Protocol):
     """
     A protocol to proxy data between two transports.
     """
@@ -738,7 +838,7 @@ async def test_Connector_connect_socket_async(
     ) as connector:
         logger.info("client socket opening")
         connector._client = fake_client
-        p = TestProtocol()
+        p = SocketTestProtocol()
 
         # Open proxy connection
         # start the proxy server
@@ -759,3 +859,72 @@ async def test_Connector_connect_socket_async(
         logger.info("client socket done")
 
         assert p.received.decode() == "world\n"
+
+
+@pytest.mark.asyncio
+async def test_Connector_connect_socket_async_invalidates_cache_on_connection_error(
+    fake_credentials: Credentials,
+) -> None:
+    connect_string = "test-project:test-region:test-instance"
+    async with Connector(
+        credentials=fake_credentials, loop=asyncio.get_running_loop()
+    ) as connector:
+        monitored_cache = MagicMock()
+        conn_info = MagicMock()
+        conn_info.create_ssl_context = AsyncMock(return_value=object())
+        conn_info.get_preferred_ip.return_value = "127.0.0.1"
+        monitored_cache.connect_info = AsyncMock(return_value=conn_info)
+
+        with (
+            patch.object(
+                connector, "_get_cache", AsyncMock(return_value=monitored_cache)
+            ),
+            patch.object(
+                connector._loop,
+                "create_connection",
+                AsyncMock(side_effect=RuntimeError("boom")),
+            ),
+            patch.object(
+                connector, "_remove_cached", AsyncMock()
+            ) as mock_remove_cached,
+        ):
+            with pytest.raises(RuntimeError) as exc_info:
+                await connector.connect_socket_async(
+                    connect_string,
+                    asyncio.Protocol,
+                    driver="asyncpg",
+                )
+
+        assert exc_info.value.args[0] == "boom"
+        mock_remove_cached.assert_awaited_once_with(connect_string, False)
+
+
+@pytest.mark.asyncio
+async def test_ConnectorSocketFactory_connect_forwards_arguments(
+    fake_credentials: Credentials,
+) -> None:
+    connect_string = "test-project:test-region:test-instance"
+    async with Connector(
+        credentials=fake_credentials, loop=asyncio.get_running_loop()
+    ) as connector:
+        protocol_fn = MagicMock()
+        with patch.object(
+            connector,
+            "connect_socket_async",
+            AsyncMock(),
+        ) as mock_connect_socket_async:
+            factory = ConnectorSocketFactory(
+                connector,
+                connect_string,
+                driver="asyncpg",
+                user="my-user",
+            )
+
+            await factory.connect(protocol_fn)
+
+        mock_connect_socket_async.assert_awaited_once_with(
+            connect_string,
+            protocol_fn,
+            driver="asyncpg",
+            user="my-user",
+        )

tests/unit/test_proxy.py

@@ -22,7 +22,10 @@
 
 import pytest
 
+from google.cloud.sql.connector.proxy import BaseProxyProtocol
+from google.cloud.sql.connector.proxy import ClientToServerProtocol
 from google.cloud.sql.connector.proxy import Proxy
+from google.cloud.sql.connector.proxy import ProxyClientConnection
 from google.cloud.sql.connector.proxy import ServerConnectionFactory
 
 
@@ -50,6 +53,70 @@ async def test_proxy_creates_folder_and_socket(short_tmpdir):
     await proxy.close()
 
 
+def test_base_proxy_protocol_flushes_cached_data_when_target_is_set():
+    protocol = BaseProxyProtocol(MagicMock())
+    target = MagicMock(spec=asyncio.Transport)
+
+    protocol.data_received(b"hello")
+    protocol.data_received(b"world")
+    protocol.set_target(target)
+
+    target.writelines.assert_called_once_with([b"hello", b"world"])
+    assert protocol._cached == []
+
+
+def test_base_proxy_protocol_forwards_eof_and_close_to_target():
+    protocol = BaseProxyProtocol(MagicMock())
+    target = MagicMock(spec=asyncio.Transport)
+    protocol.set_target(target)
+
+    protocol.eof_received()
+    protocol.connection_lost(None)
+
+    target.write_eof.assert_called_once()
+    target.close.assert_called_once()
+
+
+def test_proxy_client_connection_close_prefers_write_eof():
+    client_transport = MagicMock(spec=asyncio.Transport)
+    client_transport.is_closing.return_value = False
+    client_transport.can_write_eof.return_value = True
+
+    server_transport = MagicMock(spec=asyncio.Transport)
+    server_transport.is_closing.return_value = False
+    server_transport.can_write_eof.return_value = True
+
+    connection = ProxyClientConnection(client_transport, MagicMock())
+    connection.server_transport = server_transport
+
+    connection.close()
+
+    client_transport.write_eof.assert_called_once()
+    server_transport.write_eof.assert_called_once()
+
+
+def test_proxy_client_connection_close_falls_back_to_close():
+    client_transport = MagicMock(spec=asyncio.Transport)
+    client_transport.is_closing.return_value = False
+    client_transport.can_write_eof.return_value = False
+
+    connection = ProxyClientConnection(client_transport, MagicMock())
+
+    connection.close()
+
+    client_transport.close.assert_called_once()
+
+
+def test_client_to_server_protocol_opens_backend_connection_on_accept():
+    proxy = MagicMock()
+    protocol = ClientToServerProtocol(proxy)
+    transport = MagicMock(spec=asyncio.Transport)
+
+    protocol.connection_made(transport)
+
+    proxy._handle_client_connection.assert_called_once_with(transport, protocol)
+
+
 # A mock ServerConnectionFactory for testing purposes.
 class MockServerConnectionFactory(ServerConnectionFactory):
     def __init__(self, loop):
@@ -373,4 +440,4 @@ async def test_tcp_proxy_client_closes_connection(tcp_proxy_server):
 
     # Check that the server socket is closing
     await asyncio.sleep(0.01)
-    assert connector.server_transport.is_closing()
+    assert connector.server_transport.is_closing()