diff --git a/google/cloud/sql/connector/asyncpg.py b/google/cloud/sql/connector/asyncpg.py index 2fbc3027..191a5618 100644 --- a/google/cloud/sql/connector/asyncpg.py +++ b/google/cloud/sql/connector/asyncpg.py @@ -51,7 +51,9 @@ async def connect( 'Unable to import module "asyncpg." Please install and try again.' ) user = kwargs.pop("user") - db = kwargs.pop("db") + db = kwargs.pop("database", kwargs.pop("db", None)) + if db is None: + raise KeyError("database") passwd = kwargs.pop("password", None) return await asyncpg.connect( diff --git a/google/cloud/sql/connector/pg8000.py b/google/cloud/sql/connector/pg8000.py index 5a43ad31..89eaed41 100644 --- a/google/cloud/sql/connector/pg8000.py +++ b/google/cloud/sql/connector/pg8000.py @@ -48,7 +48,9 @@ def connect( ) user = kwargs.pop("user") - db = kwargs.pop("db") + db = kwargs.pop("database", kwargs.pop("db", None)) + if db is None: + raise KeyError("database") passwd = kwargs.pop("password", None) return pg8000.dbapi.connect( user, diff --git a/google/cloud/sql/connector/pymysql.py b/google/cloud/sql/connector/pymysql.py index c087e9bc..7ec6607a 100644 --- a/google/cloud/sql/connector/pymysql.py +++ b/google/cloud/sql/connector/pymysql.py @@ -52,6 +52,12 @@ def connect( # pop timeout as timeout arg is called 'connect_timeout' for pymysql timeout = kwargs.pop("timeout") kwargs["connect_timeout"] = kwargs.get("connect_timeout", timeout) + + # map 'db' to 'database' to avoid deprecation warning in pymysql + db = kwargs.pop("db", None) + if db is not None: + kwargs["database"] = db + # Create pymysql connection object and hand in pre-made connection conn = pymysql.Connection(host=ip_address, defer_connect=True, **kwargs) conn.connect(sock) diff --git a/google/cloud/sql/connector/pytds.py b/google/cloud/sql/connector/pytds.py index 6cc3c093..76a633bf 100644 --- a/google/cloud/sql/connector/pytds.py +++ b/google/cloud/sql/connector/pytds.py @@ -47,7 +47,7 @@ def connect(ip_address: str, sock: ssl.SSLSocket, **kwargs: Any) -> "pytds.Conne 'Unable to import module "pytds." Please install and try again.' ) - db = kwargs.pop("db", None) + db = kwargs.pop("database", kwargs.pop("db", None)) if kwargs.pop("active_directory_auth", False): if platform.system() == "Windows": diff --git a/tests/system/test_asyncpg_connection.py b/tests/system/test_asyncpg_connection.py index 00589942..72f4dabb 100644 --- a/tests/system/test_asyncpg_connection.py +++ b/tests/system/test_asyncpg_connection.py @@ -92,7 +92,7 @@ async def create_sqlalchemy_engine( "asyncpg", user=user, password=password, - db=db, + database=db, ip_type=ip_type, # can be "public", "private" or "psc" **kwargs, # additional asyncpg connection args ), @@ -155,7 +155,7 @@ async def getconn( "asyncpg", user=user, password=password, - db=db, + database=db, ip_type=ip_type, # can be "public", "private" or "psc" **kwargs, ) diff --git a/tests/system/test_asyncpg_iam_auth.py b/tests/system/test_asyncpg_iam_auth.py index 0e0c01e8..a898352a 100644 --- a/tests/system/test_asyncpg_iam_auth.py +++ b/tests/system/test_asyncpg_iam_auth.py @@ -74,7 +74,7 @@ async def create_sqlalchemy_engine( instance_connection_name, "asyncpg", user=user, - db=db, + database=db, ip_type=ip_type, # can be "public", "private" or "psc" enable_iam_auth=True, ), diff --git a/tests/system/test_connector_object.py b/tests/system/test_connector_object.py index 5b05319f..89610600 100644 --- a/tests/system/test_connector_object.py +++ b/tests/system/test_connector_object.py @@ -37,7 +37,7 @@ def getconn() -> pymysql.connections.Connection: "pymysql", user=os.environ["MYSQL_USER"], password=os.environ["MYSQL_PASS"], - db=os.environ["MYSQL_DB"], + database=os.environ["MYSQL_DB"], ip_type=os.environ.get("IP_TYPE", "public"), ) return conn @@ -142,5 +142,5 @@ def test_connector_sqlserver_iam_auth_error() -> None: "pytds", user="my-user", password="my-pass", - db="my-db", + database="my-db", ) diff --git a/tests/system/test_ip_types.py b/tests/system/test_ip_types.py index 3af49c54..3f595e00 100644 --- a/tests/system/test_ip_types.py +++ b/tests/system/test_ip_types.py @@ -37,7 +37,7 @@ def getconn() -> pymysql.connections.Connection: ip_type=ip_type, user=os.environ["MYSQL_USER"], password=os.environ["MYSQL_PASS"], - db=os.environ["MYSQL_DB"], + database=os.environ["MYSQL_DB"], ) return conn diff --git a/tests/system/test_pg8000_connection.py b/tests/system/test_pg8000_connection.py index 0fef1b96..4a2464cb 100644 --- a/tests/system/test_pg8000_connection.py +++ b/tests/system/test_pg8000_connection.py @@ -87,7 +87,7 @@ def create_sqlalchemy_engine( "pg8000", user=user, password=password, - db=db, + database=db, ip_type=ip_type, # can be "public", "private" or "psc" ), ) diff --git a/tests/system/test_pg8000_iam_auth.py b/tests/system/test_pg8000_iam_auth.py index 38ee7661..63aa8b5a 100644 --- a/tests/system/test_pg8000_iam_auth.py +++ b/tests/system/test_pg8000_iam_auth.py @@ -73,7 +73,7 @@ def create_sqlalchemy_engine( instance_connection_name, "pg8000", user=user, - db=db, + database=db, ip_type=ip_type, # can be "public", "private" or "psc" enable_iam_auth=True, ), diff --git a/tests/system/test_pymysql_connection.py b/tests/system/test_pymysql_connection.py index f5dd8a1e..604a3b95 100644 --- a/tests/system/test_pymysql_connection.py +++ b/tests/system/test_pymysql_connection.py @@ -78,7 +78,7 @@ def create_sqlalchemy_engine( "pymysql", user=user, password=password, - db=db, + database=db, ip_type=ip_type, # can be "public", "private" or "psc" ), ) diff --git a/tests/system/test_pymysql_iam_auth.py b/tests/system/test_pymysql_iam_auth.py index da676e66..8117f8fc 100644 --- a/tests/system/test_pymysql_iam_auth.py +++ b/tests/system/test_pymysql_iam_auth.py @@ -73,7 +73,7 @@ def create_sqlalchemy_engine( instance_connection_name, "pymysql", user=user, - db=db, + database=db, ip_type=ip_type, # can be "public", "private" or "psc" enable_iam_auth=True, ), diff --git a/tests/system/test_pytds_connection.py b/tests/system/test_pytds_connection.py index 198c3307..7e113e3c 100644 --- a/tests/system/test_pytds_connection.py +++ b/tests/system/test_pytds_connection.py @@ -76,7 +76,7 @@ def create_sqlalchemy_engine( "pytds", user=user, password=password, - db=db, + database=db, ip_type=ip_type, # can be "public", "private" or "psc" ), ) diff --git a/tests/unit/test_asyncpg.py b/tests/unit/test_asyncpg.py index ee1da1ec..74961ae3 100644 --- a/tests/unit/test_asyncpg.py +++ b/tests/unit/test_asyncpg.py @@ -35,3 +35,64 @@ async def test_asyncpg(mock_connect: AsyncMock, kwargs: Any) -> None: assert connection is True # verify that driver connection call would be made assert mock_connect.assert_called_once + + +@pytest.mark.asyncio +async def test_asyncpg_import_error(kwargs: Any) -> None: + """Test to verify that connect raises ImportError if asyncpg is not installed.""" + with patch.dict("sys.modules", {"asyncpg": None}): + with pytest.raises(ImportError) as excinfo: + await connect("0.0.0.0", ssl.create_default_context(), **kwargs) + assert 'Unable to import module "asyncpg."' in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_asyncpg_database_param(kwargs: Any) -> None: + """Test that asyncpg wrapper accepts both 'db' and 'database'.""" + ip_addr = "0.0.0.0" + context = ssl.create_default_context() + + # Test with 'db' + kwargs_db = kwargs.copy() + kwargs_db["db"] = "my-db-1" + with patch("asyncpg.connect", new_callable=AsyncMock) as mock_connect: + await connect(ip_addr, context, **kwargs_db) + mock_connect.assert_called_once_with( + user=kwargs["user"], + database="my-db-1", + password=kwargs.get("password"), + host=ip_addr, + port=3307, + ssl=context, + direct_tls=True, + ) + + # Test with 'database' + kwargs_database = kwargs.copy() + kwargs_database["database"] = "my-db-2" + if "db" in kwargs_database: + del kwargs_database["db"] + with patch("asyncpg.connect", new_callable=AsyncMock) as mock_connect: + await connect(ip_addr, context, **kwargs_database) + mock_connect.assert_called_once_with( + user=kwargs["user"], + database="my-db-2", + password=kwargs.get("password"), + host=ip_addr, + port=3307, + ssl=context, + direct_tls=True, + ) + + +@pytest.mark.asyncio +async def test_asyncpg_database_missing(kwargs: Any) -> None: + """Test that asyncpg wrapper raises KeyError if database is missing.""" + if "db" in kwargs: + del kwargs["db"] + if "database" in kwargs: + del kwargs["database"] + with pytest.raises(KeyError, match="database"): + await connect("0.0.0.0", ssl.create_default_context(), **kwargs) + + diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index e9f7eb88..bdb3e568 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -15,14 +15,19 @@ from __future__ import annotations import datetime +from unittest.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import patch from aiohttp import ClientResponseError from aioresponses import aioresponses from google.auth.credentials import Credentials +from google.auth.credentials import TokenState from mocks import FakeCredentials import pytest from google.cloud.sql.connector.client import CloudSQLClient +from google.cloud.sql.connector.client import DEFAULT_SERVICE_ENDPOINT from google.cloud.sql.connector.utils import generate_keys from google.cloud.sql.connector.version import __version__ as version @@ -323,3 +328,228 @@ async def test_get_metadata_multiple_psc_dns_sorted(fake_client: CloudSQLClient) fake_client.instance.legacy_dns_name = False fake_client.instance.dns_names = ["abcde.12345.us-central1.sql.goog"] + +async def test_CloudSQLClient_init_default_endpoint( + fake_credentials: FakeCredentials, +) -> None: + """Test that CloudSQLClient uses default endpoint if None is passed.""" + client = CloudSQLClient(None, "my-quota-project", fake_credentials) + assert client._sqladmin_api_endpoint == DEFAULT_SERVICE_ENDPOINT + await client.close() + + +async def test_get_metadata_retry_50x(fake_credentials: Credentials) -> None: + """Test that _get_metadata retries on 5xx errors.""" + client = CloudSQLClient( + sqladmin_api_endpoint="https://sqladmin.googleapis.com", + quota_project=None, + credentials=fake_credentials, + ) + get_url = "https://sqladmin.googleapis.com/sql/v1beta4/projects/my-project/instances/my-instance/connectSettings" + + resp_body = { + "ipAddresses": [{"type": "PRIMARY", "ipAddress": "127.0.0.1"}], + "region": "my-region", + "databaseVersion": "POSTGRES_15", + "serverCaCert": {"cert": "ca-cert"}, + } + + with aioresponses() as mocked: + # First call returns 500 + mocked.get(get_url, status=500) + # Second call returns 200 + mocked.get(get_url, status=200, payload=resp_body) + + # We need to mock sleep in retry_50x to make it fast + with patch( + "google.cloud.sql.connector.refresh_utils.asyncio.sleep", AsyncMock() + ): + resp = await client._get_metadata("my-project", "my-region", "my-instance") + + assert resp["database_version"] == "POSTGRES_15" + assert resp["ip_addresses"] == {"PRIMARY": ["127.0.0.1"]} + await client.close() + + +async def test_get_ephemeral_retry_50x(fake_credentials: Credentials) -> None: + """Test that _get_ephemeral retries on 5xx errors.""" + client = CloudSQLClient( + sqladmin_api_endpoint="https://sqladmin.googleapis.com", + quota_project=None, + credentials=fake_credentials, + ) + post_url = "https://sqladmin.googleapis.com/sql/v1beta4/projects/my-project/instances/my-instance:generateEphemeralCert" + + resp_body = { + "ephemeralCert": {"cert": "ephemeral-cert"}, + } + + mock_x509 = MagicMock() + mock_x509.not_valid_after_utc = datetime.datetime.now( + datetime.timezone.utc + ) + datetime.timedelta(hours=1) + + with aioresponses() as mocked, patch( + "google.cloud.sql.connector.client.load_pem_x509_certificate", + return_value=mock_x509, + ): + # First call returns 500 + mocked.post(post_url, status=500) + # Second call returns 200 + mocked.post(post_url, status=200, payload=resp_body) + + # We need to mock sleep in retry_50x to make it fast + with patch( + "google.cloud.sql.connector.refresh_utils.asyncio.sleep", AsyncMock() + ): + cert, expiration = await client._get_ephemeral( + "my-project", "my-instance", "pub-key" + ) + + assert cert == "ephemeral-cert" + assert expiration == mock_x509.not_valid_after_utc + await client.close() + + +async def test_get_metadata_region_mismatch(fake_credentials: Credentials) -> None: + """Test that _get_metadata raises ValueError if region mismatched.""" + client = CloudSQLClient( + sqladmin_api_endpoint="https://sqladmin.googleapis.com", + quota_project=None, + credentials=fake_credentials, + ) + get_url = "https://sqladmin.googleapis.com/sql/v1beta4/projects/my-project/instances/my-instance/connectSettings" + + resp_body = { + "ipAddresses": [{"type": "PRIMARY", "ipAddress": "127.0.0.1"}], + "region": "wrong-region", + "databaseVersion": "POSTGRES_15", + "serverCaCert": {"cert": "ca-cert"}, + } + + with aioresponses() as mocked: + mocked.get(get_url, status=200, payload=resp_body) + with pytest.raises(ValueError) as exc_info: + await client._get_metadata("my-project", "my-region", "my-instance") + assert "Provided region was mismatched" in str(exc_info.value) + await client.close() + + +async def test_resolve_connect_settings_success(fake_credentials: Credentials) -> None: + """Test resolve_connect_settings returns successfully.""" + client = CloudSQLClient( + sqladmin_api_endpoint="https://sqladmin.googleapis.com", + quota_project=None, + credentials=fake_credentials, + ) + get_url = "https://sqladmin.googleapis.com/sql/v1beta4/locations/my-region/dns/my-dns:resolveConnectSettings" + resp_body = {"connectionName": "my-project:my-region:my-instance"} + + with aioresponses() as mocked: + mocked.get(get_url, status=200, payload=resp_body) + resp = await client.resolve_connect_settings("my-dns", "my-region") + assert resp == resp_body + await client.close() + + +async def test_resolve_connect_settings_error(fake_credentials: Credentials) -> None: + """Test resolve_connect_settings raises error on failure.""" + client = CloudSQLClient( + sqladmin_api_endpoint="https://sqladmin.googleapis.com", + quota_project=None, + credentials=fake_credentials, + ) + get_url = "https://sqladmin.googleapis.com/sql/v1beta4/locations/my-region/dns/my-dns:resolveConnectSettings" + resp_body = { + "error": { + "code": 404, + "message": "DNS name not found", + } + } + + with aioresponses() as mocked: + mocked.get(get_url, status=404, payload=resp_body) + with pytest.raises(ClientResponseError) as exc_info: + await client.resolve_connect_settings("my-dns", "my-region") + assert exc_info.value.status == 404 + assert exc_info.value.message == "DNS name not found" + await client.close() + + +async def test_resolve_connect_settings_retry_50x(fake_credentials: Credentials) -> None: + """Test that resolve_connect_settings retries on 5xx errors.""" + client = CloudSQLClient( + sqladmin_api_endpoint="https://sqladmin.googleapis.com", + quota_project=None, + credentials=fake_credentials, + ) + get_url = "https://sqladmin.googleapis.com/sql/v1beta4/locations/my-region/dns/my-dns:resolveConnectSettings" + resp_body = {"connectionName": "my-project:my-region:my-instance"} + + with aioresponses() as mocked: + # First call returns 500 + mocked.get(get_url, status=500) + # Second call returns 200 + mocked.get(get_url, status=200, payload=resp_body) + + with patch( + "google.cloud.sql.connector.refresh_utils.asyncio.sleep", AsyncMock() + ): + resp = await client.resolve_connect_settings("my-dns", "my-region") + + assert resp == resp_body + await client.close() + + +async def test_resolve_connect_settings_token_refresh(fake_credentials: FakeCredentials) -> None: + """Test that resolve_connect_settings refreshes token if it is not FRESH.""" + fake_credentials.token = "expired-token" + fake_credentials.expiry = datetime.datetime.now( + datetime.timezone.utc + ) - datetime.timedelta(minutes=10) + assert fake_credentials.token_state == TokenState.INVALID + + client = CloudSQLClient( + sqladmin_api_endpoint="https://sqladmin.googleapis.com", + quota_project=None, + credentials=fake_credentials, + ) + get_url = "https://sqladmin.googleapis.com/sql/v1beta4/locations/my-region/dns/my-dns:resolveConnectSettings" + resp_body = {"connectionName": "my-project:my-region:my-instance"} + + with aioresponses() as mocked: + mocked.get(get_url, status=200, payload=resp_body) + resp = await client.resolve_connect_settings("my-dns", "my-region") + + assert resp == resp_body + assert fake_credentials.token == "12345" + await client.close() + + +async def test_resolve_connect_settings_error_parsing_json( + fake_credentials: Credentials, +) -> None: + """Test that aiohttp default error messages are raised when resolve_connect_settings gets a bad JSON response.""" + client = CloudSQLClient( + sqladmin_api_endpoint="https://sqladmin.googleapis.com", + quota_project=None, + credentials=fake_credentials, + ) + get_url = "https://sqladmin.googleapis.com/sql/v1beta4/locations/my-region/dns/my-dns:resolveConnectSettings" + resp_body = ["error"] + with aioresponses() as mocked: + mocked.get( + get_url, + status=403, + payload=resp_body, + repeat=True, + ) + with pytest.raises(ClientResponseError) as exc_info: + await client.resolve_connect_settings("my-dns", "my-region") + assert exc_info.value.status == 403 + assert exc_info.value.message == "Forbidden" + await client.close() + + + + diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py index 8ea8c88c..edc2bf3e 100644 --- a/tests/unit/test_connector.py +++ b/tests/unit/test_connector.py @@ -18,6 +18,8 @@ import asyncio import os from threading import Thread +from unittest.mock import AsyncMock +from unittest.mock import MagicMock from unittest.mock import patch from aiohttp import ClientResponseError @@ -29,11 +31,13 @@ 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.enums import RefreshStrategy from google.cloud.sql.connector.exceptions import ClosedConnectorError from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError from google.cloud.sql.connector.exceptions import ConnectorLoopError from google.cloud.sql.connector.exceptions import IncompatibleDriverError from google.cloud.sql.connector.instance import RefreshAheadCache +from google.cloud.sql.connector.monitored_cache import MonitoredCache from google.cloud.sql.connector.resolver import DnsResolver @@ -101,6 +105,13 @@ async def test_connect_incompatible_driver_error( ) +def test_connector_invalid_refresh_strategy() -> None: + """Test that initializing Connector with invalid refresh_strategy raises ValueError.""" + with pytest.raises(ValueError) as exc_info: + Connector(refresh_strategy="INVALID") + assert "Incorrect value for refresh_strategy, got 'INVALID'" in str(exc_info.value) + + def test_connect_with_unsupported_driver(fake_credentials: Credentials) -> None: with Connector(credentials=fake_credentials) as connector: # try to connect using unsupported driver, should raise KeyError @@ -722,3 +733,304 @@ async def test_Connector_connect_async_custom_dns_resolver_no_fallback_psc_to_pr fake_client.instance.psc_enabled = False +def test_Connector_init_lazy_with_loop(fake_credentials: Credentials) -> None: + """Test Connector initialization with custom loop and LAZY refresh strategy.""" + loop = asyncio.new_event_loop() + try: + connector = Connector( + credentials=fake_credentials, + loop=loop, + refresh_strategy=RefreshStrategy.LAZY, + ) + assert connector._loop == loop + assert connector._refresh_strategy == RefreshStrategy.LAZY + assert connector._keys is None + finally: + loop.close() + + +def test_Connector_init_quota_project(fake_credentials: Credentials) -> None: + """Test Connector initialization with quota_project argument.""" + connector = Connector(credentials=fake_credentials, quota_project="my-quota-project") + assert connector._quota_project == "my-quota-project" + + +def test_Connector_init_sqladmin_endpoint(fake_credentials: Credentials) -> None: + """Test Connector initialization with sqladmin_api_endpoint argument.""" + connector = Connector( + credentials=fake_credentials, + sqladmin_api_endpoint="https://my-custom-endpoint.com", + ) + assert connector._sqladmin_api_endpoint == "https://my-custom-endpoint.com" + + +@pytest.mark.asyncio +async def test_Connector_connect_async_sync_driver( + fake_credentials: Credentials, fake_client: CloudSQLClient +) -> None: + """Test that Connector.connect_async works with a sync driver (pg8000).""" + async with Connector( + credentials=fake_credentials, loop=asyncio.get_running_loop() + ) as connector: + connector._client = fake_client + + # Mock socket.create_connection to avoid real network call + # Mock ssl.SSLContext.wrap_socket to return a mock socket + mock_sock = MagicMock() + with patch( + "google.cloud.sql.connector.connector.socket.create_connection", + return_value=mock_sock, + ), patch( + "ssl.SSLContext.wrap_socket", return_value=mock_sock + ), patch( + "google.cloud.sql.connector.pg8000.connect" + ) as mock_connect: + + mock_connect.return_value = True + + connection = await connector.connect_async( + "test-project:test-region:test-instance", + "pg8000", + user="my-user", + password="my-pass", + db="my-db", + ) + assert connection is True + mock_connect.assert_called_once() + + + +@pytest.mark.asyncio +async def test_Connector_connect_async_iam_username_truncation( + fake_credentials: Credentials, fake_client: CloudSQLClient +) -> None: + """Test that IAM username is truncated and warning is logged.""" + async with Connector( + credentials=fake_credentials, loop=asyncio.get_running_loop() + ) as connector: + connector._client = fake_client + + # Modify database version to POSTGRES to trigger truncation for postgres SA + fake_client.instance.db_version = "POSTGRES_15" + + with patch("google.cloud.sql.connector.asyncpg.connect") as mock_connect, patch( + "google.cloud.sql.connector.connector.logger" + ) as mock_logger: + mock_connect.return_value = True + + # Use a username that needs truncation (ends with .gserviceaccount.com) + long_user = "my-sa@my-project.iam.gserviceaccount.com" + expected_user = "my-sa@my-project.iam" + + await connector.connect_async( + "test-project:test-region:test-instance", + "asyncpg", + user=long_user, + password="my-pass", + db="my-db", + enable_iam_auth=True, + ) + + # Verify truncated user was passed to driver + _, kwargs = mock_connect.call_args + assert kwargs["user"] == expected_user + + # Verify truncation warning was logged + mock_logger.debug.assert_any_call( + f"['test-project:test-region:test-instance']: Truncated IAM database username from {long_user} to {expected_user}" + ) + + +@pytest.mark.asyncio +async def test_Connector_connect_async_lazy_refresh( + fake_credentials: Credentials, fake_client: CloudSQLClient +) -> None: + """Test that Connector.connect_async works with LAZY refresh strategy.""" + async with Connector( + credentials=fake_credentials, + loop=asyncio.get_running_loop(), + refresh_strategy=RefreshStrategy.LAZY, + ) as connector: + connector._client = fake_client + + with patch("google.cloud.sql.connector.asyncpg.connect") as mock_connect: + mock_connect.return_value = True + + # self._keys should be None initially for LAZY + assert connector._keys is None + + connection = await connector.connect_async( + "test-project:test-region:test-instance", + "asyncpg", + user="my-user", + password="my-pass", + db="my-db", + ) + assert connection is True + # self._keys should now be initialized + assert connector._keys is not None + mock_connect.assert_called_once() + + +@pytest.mark.asyncio +async def test_Connector_connect_async_custom_dns_resolver_exception( + fake_credentials: Credentials, fake_client: CloudSQLClient +) -> None: + """Test that Connector.connect_async falls back if DNS resolution raises exception.""" + with patch( + "google.cloud.sql.connector.resolver.DnsResolver.resolve_a_record", + side_effect=Exception("DNS Resolution Failed"), + ), patch( + "google.cloud.sql.connector.resolver.DnsResolver.resolve" + ) as mock_resolve: + conn_name_with_domain = ConnectionName( + "test-project", "test-region", "test-instance", "db.example.com" + ) + mock_resolve.return_value = conn_name_with_domain + + async with Connector( + credentials=fake_credentials, + loop=asyncio.get_running_loop(), + resolver=DnsResolver, + ) as connector: + connector._client = fake_client + + original_ips = fake_client.instance.ip_addrs + fake_client.instance.ip_addrs = {"PRIMARY": "5.6.7.8"} + + try: + with patch( + "google.cloud.sql.connector.asyncpg.connect" + ) as mock_connect: + mock_connect.return_value = True + + connection = await connector.connect_async( + "db.example.com", + "asyncpg", + user="my-user", + password="my-pass", + db="my-db", + ) + + # Verify mock_connect was called with metadata IP "5.6.7.8" + args, _ = mock_connect.call_args + assert args[0] == "5.6.7.8" + assert connection is True + finally: + fake_client.instance.ip_addrs = original_ips + + +@pytest.mark.asyncio +async def test_Connector_connect_async_sync_driver_ssl_error( + fake_credentials: Credentials, fake_client: CloudSQLClient +) -> None: + """Test that Connector.connect_async closes raw socket if SSL wrap fails.""" + async with Connector( + credentials=fake_credentials, loop=asyncio.get_running_loop() + ) as connector: + connector._client = fake_client + + mock_raw_sock = MagicMock() + with patch( + "google.cloud.sql.connector.connector.socket.create_connection", + return_value=mock_raw_sock, + ), patch( + "ssl.SSLContext.wrap_socket", side_effect=Exception("SSL wrap failed") + ), patch( + "google.cloud.sql.connector.pg8000.connect" + ) as mock_connect: + + with pytest.raises(Exception, match="SSL wrap failed"): + await connector.connect_async( + "test-project:test-region:test-instance", + "pg8000", + user="my-user", + password="my-pass", + db="my-db", + ) + + mock_raw_sock.close.assert_called_once() + mock_connect.assert_not_called() + + +@pytest.mark.asyncio +async def test_Connector_connect_async_sync_driver_domain_name( + fake_credentials: Credentials, fake_client: CloudSQLClient +) -> None: + """Test that Connector.connect_async stores socket in monitored_cache if domain name is used.""" + with patch( + "google.cloud.sql.connector.resolver.DnsResolver.resolve_a_record", + return_value=["1.2.3.4"], + ), patch( + "google.cloud.sql.connector.resolver.DnsResolver.resolve" + ) as mock_resolve: + conn_name_with_domain = ConnectionName( + "test-project", "test-region", "test-instance", "db.example.com" + ) + mock_resolve.return_value = conn_name_with_domain + + async with Connector( + credentials=fake_credentials, + loop=asyncio.get_running_loop(), + resolver=DnsResolver, + ) as connector: + connector._client = fake_client + + mock_sock = MagicMock() + with patch( + "google.cloud.sql.connector.connector.socket.create_connection", + return_value=mock_sock, + ), patch( + "ssl.SSLContext.wrap_socket", return_value=mock_sock + ), patch( + "google.cloud.sql.connector.pg8000.connect" + ) as mock_connect: + + mock_connect.return_value = True + + connection = await connector.connect_async( + "db.example.com", + "pg8000", + user="my-user", + password="my-pass", + db="my-db", + ) + assert connection is True + + # Get monitored cache from connector + monitored_cache = connector._cache[(str(conn_name_with_domain), False)] + # Verify mock_sock was appended to monitored_cache.sockets + assert mock_sock in monitored_cache.sockets + + +@pytest.mark.asyncio +async def test_Connector_connect_async_connection_error_triggers_force_refresh( + fake_credentials: Credentials, fake_client: CloudSQLClient +) -> None: + """Test that connection error triggers force_refresh on monitored cache.""" + async with Connector( + credentials=fake_credentials, loop=asyncio.get_running_loop() + ) as connector: + connector._client = fake_client + + with patch( + "google.cloud.sql.connector.asyncpg.connect", + side_effect=Exception("Connection Refused"), + ), patch.object( + MonitoredCache, "force_refresh", AsyncMock() + ) as mock_force_refresh: + + with pytest.raises(Exception, match="Connection Refused"): + await connector.connect_async( + "test-project:test-region:test-instance", + "asyncpg", + user="my-user", + password="my-pass", + db="my-db", + ) + + mock_force_refresh.assert_called_once() + + + + diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py index 3dbba59b..ca87370b 100644 --- a/tests/unit/test_instance.py +++ b/tests/unit/test_instance.py @@ -16,6 +16,8 @@ import asyncio import datetime +import ssl +from unittest.mock import AsyncMock from unittest.mock import patch import mocks @@ -27,6 +29,7 @@ from google.cloud.sql.connector.connection_name import ConnectionName from google.cloud.sql.connector.exceptions import AutoIAMAuthNotSupported from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError +from google.cloud.sql.connector.exceptions import TLSVersionError from google.cloud.sql.connector.instance import RefreshAheadCache from google.cloud.sql.connector.rate_limiter import AsyncRateLimiter from google.cloud.sql.connector.refresh_utils import _is_valid @@ -166,6 +169,43 @@ async def test_force_refresh_cancels_pending_refresh( assert isinstance(await cache._current, ConnectionInfo) +@pytest.mark.asyncio +async def test_force_refresh_replaces_invalid_current( + cache: RefreshAheadCache, + test_rate_limiter: AsyncRateLimiter, +) -> None: + """Test that force_refresh replaces current task with next task if current is invalid.""" + cache._refresh_rate_limiter = test_rate_limiter + # make sure initial refresh is finished + await cache._current + + # Create an expired ConnectionInfo + expired_info = ConnectionInfo( + cache._conn_name, + "cert", + "ca", + b"key", + {}, + "POSTGRES", + datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10), + ) + + # Create a task that returns the expired ConnectionInfo + async def get_expired() -> ConnectionInfo: + return expired_info + + cache._current = asyncio.create_task(get_expired()) + await cache._current # wait for it to complete so _is_valid can read it + + assert await _is_valid(cache._current) is False + + next_refresh = cache._next + await cache.force_refresh() + + assert cache._current == cache._next + assert cache._current != next_refresh + + @pytest.mark.asyncio async def test_RefreshAheadCache_close(cache: RefreshAheadCache) -> None: """ @@ -293,3 +333,61 @@ async def test_ConnectionInfo_caches_sslcontext() -> None: # calling create_ssl_context should no-op with an existing 'context' await info.create_ssl_context() assert info.context == "context" + + +@pytest.mark.asyncio +async def test_ConnectionInfo_create_ssl_context_no_tls1_3_error() -> None: + """Test that create_ssl_context raises TLSVersionError if TLSv1.3 is not supported and IAM is enabled.""" + info = ConnectionInfo( + ConnectionName("p", "r", "i"), + "cert", + "ca", + b"key", + {}, + "POSTGRES", + datetime.datetime.now(datetime.timezone.utc), + ) + with patch("google.cloud.sql.connector.connection_info.ssl.HAS_TLSv1_3", False): + with pytest.raises(TLSVersionError) as exc_info: + await info.create_ssl_context(enable_iam_auth=True) + assert ( + "does not support TLSv1.3, which is required to use IAM Authentication" + in str(exc_info.value) + ) + + +@pytest.mark.asyncio +async def test_ConnectionInfo_create_ssl_context_no_tls1_3_warning() -> None: + """Test that create_ssl_context falls back to TLSv1.2 with warning if TLSv1.3 is not supported.""" + info = ConnectionInfo( + ConnectionName("p", "r", "i"), + "cert", + "ca", + b"key", + {}, + "POSTGRES", + datetime.datetime.now(datetime.timezone.utc), + ) + + with patch( + "google.cloud.sql.connector.connection_info.ssl.HAS_TLSv1_3", False + ), patch( + "google.cloud.sql.connector.connection_info.logger" + ) as mock_logger, patch( + "google.cloud.sql.connector.connection_info.ssl.SSLContext.load_cert_chain" + ), patch( + "google.cloud.sql.connector.connection_info.ssl.SSLContext.load_verify_locations" + ), patch( + "google.cloud.sql.connector.connection_info.write_to_file", + AsyncMock(return_value=("cert", "ca", "key")), + ): + + context = await info.create_ssl_context(enable_iam_auth=False) + + assert context.minimum_version == ssl.TLSVersion.TLSv1_2 + mock_logger.warning.assert_called_once() + assert ( + "TLSv1.3 is not supported with your version of OpenSSL" + in mock_logger.warning.call_args[0][0] + ) + diff --git a/tests/unit/test_lazy.py b/tests/unit/test_lazy.py index c6eef750..4b03ec05 100644 --- a/tests/unit/test_lazy.py +++ b/tests/unit/test_lazy.py @@ -13,6 +13,9 @@ # limitations under the License. import asyncio +from unittest.mock import AsyncMock + +import pytest from google.cloud.sql.connector.client import CloudSQLClient from google.cloud.sql.connector.connection_info import ConnectionInfo @@ -84,3 +87,27 @@ async def test_LazyRefreshCache_force_refresh(fake_client: CloudSQLClient) -> No assert conn_info2 != conn_info assert cache._cached == conn_info2 await cache.close() + + +async def test_LazyRefreshCache_connect_info_error( + fake_client: CloudSQLClient, +) -> None: + """ + Test that LazyRefreshCache.connect_info propagates exceptions. + """ + keys = asyncio.create_task(generate_keys()) + cache = LazyRefreshCache( + ConnectionName("test-project", "test-region", "test-instance"), + client=fake_client, + keys=keys, + enable_iam_auth=False, + ) + + # Mock get_connection_info to raise an exception + fake_client.get_connection_info = AsyncMock(side_effect=Exception("Test Exception")) + + with pytest.raises(Exception, match="Test Exception"): + await cache.connect_info() + + await cache.close() + diff --git a/tests/unit/test_monitored_cache.py b/tests/unit/test_monitored_cache.py index 1460ba1d..e5d87d0d 100644 --- a/tests/unit/test_monitored_cache.py +++ b/tests/unit/test_monitored_cache.py @@ -15,6 +15,7 @@ import asyncio import socket import ssl +from unittest.mock import AsyncMock from unittest.mock import patch import dns.message @@ -28,6 +29,7 @@ from google.cloud.sql.connector.exceptions import CacheClosedError from google.cloud.sql.connector.lazy import LazyRefreshCache from google.cloud.sql.connector.monitored_cache import MonitoredCache +from google.cloud.sql.connector.monitored_cache import ticker from google.cloud.sql.connector.resolver import DefaultResolver from google.cloud.sql.connector.resolver import DnsResolver from google.cloud.sql.connector.utils import generate_keys @@ -238,3 +240,119 @@ async def test_MonitoredCache_purge_closed_sockets( # call _purge_closed_sockets and verify socket is removed monitored_cache._purge_closed_sockets() assert len(monitored_cache.sockets) == 0 + + +async def test_MonitoredCache_check_domain_name_error( + fake_client: CloudSQLClient, +) -> None: + """ + Test that MonitoredCache._check_domain_name handles exceptions gracefully. + """ + conn_name = ConnectionName( + "my-project", "my-region", "my-instance", "db.example.com" + ) + cache = LazyRefreshCache( + conn_name, + client=fake_client, + keys=asyncio.create_task(generate_keys()), + enable_iam_auth=False, + ) + resolver = DnsResolver() + # Mock resolver.resolve to raise exception + resolver.resolve = AsyncMock(side_effect=Exception("DNS Error")) + + monitored_cache = MonitoredCache(cache, 0, resolver) + # verify cache is not closed initially + assert monitored_cache.closed is False + # call _check_domain_name, should catch exception and not raise/close + await monitored_cache._check_domain_name() + assert monitored_cache.closed is False + + +async def test_MonitoredCache_force_refresh_closed( + fake_client: CloudSQLClient, +) -> None: + """ + Test that MonitoredCache.force_refresh returns early if cache is closed. + """ + conn_name = ConnectionName("test-project", "test-region", "test-instance") + cache = LazyRefreshCache( + conn_name, + client=fake_client, + keys=asyncio.create_task(generate_keys()), + enable_iam_auth=False, + ) + monitored_cache = MonitoredCache(cache, 0, DefaultResolver()) + await monitored_cache.close() + assert monitored_cache.closed is True + + # Patch cache.force_refresh to verify it is NOT called + with patch.object(cache, "force_refresh", AsyncMock()) as mock_force_refresh: + await monitored_cache.force_refresh() + mock_force_refresh.assert_not_called() + + +async def test_MonitoredCache_close_twice(fake_client: CloudSQLClient) -> None: + """ + Test that MonitoredCache.close returns early if already closed. + """ + conn_name = ConnectionName("test-project", "test-region", "test-instance") + cache = LazyRefreshCache( + conn_name, + client=fake_client, + keys=asyncio.create_task(generate_keys()), + enable_iam_auth=False, + ) + monitored_cache = MonitoredCache(cache, 0, DefaultResolver()) + + async def mock_close() -> None: + cache._closed = True + + # Patch cache.close to verify it is called only once + with patch.object( + cache, "close", AsyncMock(side_effect=mock_close) + ) as mock_cache_close: + await monitored_cache.close() + assert monitored_cache.closed is True + mock_cache_close.assert_called_once() + + # call close again + await monitored_cache.close() + # should still be called only once + mock_cache_close.assert_called_once() + + +async def test_ticker() -> None: + """Test ticker function schedules calls.""" + mock_func = AsyncMock() + # Run ticker with 1s interval + task = asyncio.create_task(ticker(1, mock_func)) + # Wait for 1.5s to allow it to fire once + await asyncio.sleep(1.5) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + mock_func.assert_called_once() + + +async def test_MonitoredCache_force_refresh(fake_client: CloudSQLClient) -> None: + """ + Test that MonitoredCache.force_refresh calls cache.force_refresh. + """ + conn_name = ConnectionName("test-project", "test-region", "test-instance") + cache = LazyRefreshCache( + conn_name, + client=fake_client, + keys=asyncio.create_task(generate_keys()), + enable_iam_auth=False, + ) + monitored_cache = MonitoredCache(cache, 0, DefaultResolver()) + + # Patch cache.force_refresh to verify it is called + with patch.object(cache, "force_refresh", AsyncMock()) as mock_force_refresh: + await monitored_cache.force_refresh() + mock_force_refresh.assert_called_once() + + diff --git a/tests/unit/test_pg8000.py b/tests/unit/test_pg8000.py index c291d6f8..4ebfdc4f 100644 --- a/tests/unit/test_pg8000.py +++ b/tests/unit/test_pg8000.py @@ -17,6 +17,7 @@ import socket import ssl from typing import Any +from unittest.mock import MagicMock from unittest.mock import patch import pytest @@ -38,3 +39,55 @@ async def test_pg8000(context: ssl.SSLContext, kwargs: Any) -> None: assert connection is True # verify that driver connection call would be made assert mock_connect.assert_called_once + + +def test_pg8000_import_error(kwargs: Any) -> None: + """Test to verify that connect raises ImportError if pg8000 is not installed.""" + with patch.dict("sys.modules", {"pg8000": None}): + with pytest.raises(ImportError) as excinfo: + connect("0.0.0.0", None, **kwargs) + assert 'Unable to import module "pg8000."' in str(excinfo.value) + + +def test_pg8000_database_param(kwargs: Any) -> None: + """Test that pg8000 wrapper accepts both 'db' and 'database'.""" + ip_addr = "127.0.0.1" + sock = MagicMock(spec=ssl.SSLSocket) + + # Test with 'db' + kwargs_db = kwargs.copy() + kwargs_db["db"] = "my-db-1" + with patch("pg8000.dbapi.connect") as mock_connect: + connect(ip_addr, sock, **kwargs_db) + mock_connect.assert_called_once_with( + kwargs["user"], + database="my-db-1", + password=kwargs.get("password"), + sock=sock, + ) + + # Test with 'database' + kwargs_database = kwargs.copy() + kwargs_database["database"] = "my-db-2" + if "db" in kwargs_database: + del kwargs_database["db"] + with patch("pg8000.dbapi.connect") as mock_connect: + connect(ip_addr, sock, **kwargs_database) + mock_connect.assert_called_once_with( + kwargs["user"], + database="my-db-2", + password=kwargs.get("password"), + sock=sock, + ) + + +def test_pg8000_database_missing(kwargs: Any) -> None: + """Test that pg8000 wrapper raises KeyError if database is missing.""" + if "db" in kwargs: + del kwargs["db"] + if "database" in kwargs: + del kwargs["database"] + with pytest.raises(KeyError, match="database"): + connect("0.0.0.0", None, **kwargs) + + diff --git a/tests/unit/test_pymysql.py b/tests/unit/test_pymysql.py index ef4b23b4..bd54a515 100644 --- a/tests/unit/test_pymysql.py +++ b/tests/unit/test_pymysql.py @@ -17,6 +17,7 @@ import socket import ssl from typing import Any +from unittest.mock import MagicMock from unittest.mock import patch import pytest @@ -47,3 +48,41 @@ async def test_pymysql(context: ssl.SSLContext, kwargs: Any) -> None: pymysql_connect(ip_addr, sock, **kwargs) # verify that driver connection call would be made assert mock_connect.assert_called_once + + +def test_pymysql_import_error(kwargs: Any) -> None: + """Test to verify that connect raises ImportError if pymysql is not installed.""" + kwargs["timeout"] = 30 + with patch.dict("sys.modules", {"pymysql": None}): + with pytest.raises(ImportError) as excinfo: + pymysql_connect("0.0.0.0", None, **kwargs) + assert 'Unable to import module "pymysql."' in str(excinfo.value) + + +def test_pymysql_database_param(kwargs: Any) -> None: + """Test that pymysql wrapper accepts both 'db' and 'database' and maps 'db' to 'database'.""" + ip_addr = "127.0.0.1" + sock = MagicMock(spec=ssl.SSLSocket) + kwargs["timeout"] = 30 + + # Test with 'db' + kwargs_db = kwargs.copy() + kwargs_db["db"] = "my-db-1" + with patch("pymysql.Connection") as mock_conn: + pymysql_connect(ip_addr, sock, **kwargs_db) + _, call_kwargs = mock_conn.call_args + assert call_kwargs["database"] == "my-db-1" + assert "db" not in call_kwargs + + # Test with 'database' + kwargs_database = kwargs.copy() + kwargs_database["database"] = "my-db-2" + if "db" in kwargs_database: + del kwargs_database["db"] + with patch("pymysql.Connection") as mock_conn: + pymysql_connect(ip_addr, sock, **kwargs_database) + _, call_kwargs = mock_conn.call_args + assert call_kwargs["database"] == "my-db-2" + assert "db" not in call_kwargs + + diff --git a/tests/unit/test_pytds.py b/tests/unit/test_pytds.py index 0d9ca98c..161fa62f 100644 --- a/tests/unit/test_pytds.py +++ b/tests/unit/test_pytds.py @@ -18,6 +18,7 @@ import socket import ssl from typing import Any +from unittest.mock import MagicMock from unittest.mock import patch import pytest @@ -99,3 +100,48 @@ async def test_pytds_windows_active_directory_auth( assert mock_login.assert_called_once assert connection is True assert mock_connect.assert_called_once + + + +def test_pytds_import_error(kwargs: Any) -> None: + """Test to verify that connect raises ImportError if pytds is not installed.""" + with patch.dict("sys.modules", {"pytds": None}): + with pytest.raises(ImportError) as excinfo: + connect("0.0.0.0", None, **kwargs) + assert 'Unable to import module "pytds."' in str(excinfo.value) + + +def test_pytds_database_param(kwargs: Any) -> None: + """Test that pytds wrapper accepts both 'db' and 'database'.""" + ip_addr = "127.0.0.1" + sock = MagicMock(spec=ssl.SSLSocket) + + # Test with 'db' + kwargs_db = kwargs.copy() + kwargs_db["db"] = "my-db-1" + with patch("pytds.connect") as mock_connect: + connect(ip_addr, sock, **kwargs_db) + mock_connect.assert_called_once_with( + ip_addr, + database="my-db-1", + user=kwargs["user"], + password=kwargs["password"], + sock=sock, + ) + + # Test with 'database' + kwargs_database = kwargs.copy() + kwargs_database["database"] = "my-db-2" + if "db" in kwargs_database: + del kwargs_database["db"] + with patch("pytds.connect") as mock_connect: + connect(ip_addr, sock, **kwargs_database) + mock_connect.assert_called_once_with( + ip_addr, + database="my-db-2", + user=kwargs["user"], + password=kwargs["password"], + sock=sock, + ) + + diff --git a/tests/unit/test_refresh_utils.py b/tests/unit/test_refresh_utils.py index 0db71f18..e003b3eb 100644 --- a/tests/unit/test_refresh_utils.py +++ b/tests/unit/test_refresh_utils.py @@ -18,12 +18,14 @@ import asyncio import datetime +from typing import Any from unittest.mock import Mock from unittest.mock import patch from conftest import SCOPES # type: ignore import google.auth from google.auth.credentials import Credentials +from google.auth.credentials import Scoped from google.auth.credentials import TokenState import google.oauth2.credentials import pytest @@ -35,6 +37,25 @@ from google.cloud.sql.connector.refresh_utils import retry_50x +class DummyScopedCredentials(Credentials, Scoped): + def __init__(self, scopes: list[str] | None = None) -> None: + self._scopes = scopes + + @property + def scopes(self) -> list[str] | None: + return self._scopes + + def with_scopes(self, scopes: list[str]) -> DummyScopedCredentials: + return DummyScopedCredentials(scopes=scopes) + + def refresh(self, request: Any) -> None: + pass + + @property + def requires_scopes(self) -> bool: + return True + + @pytest.fixture def credentials() -> Credentials: credentials = Mock(spec=Credentials) @@ -101,6 +122,19 @@ def test_downscope_credentials_user() -> None: assert credentials != creds +def test_downscope_credentials_scoped() -> None: + """ + Test _downscope_credentials with a Scoped credential. + """ + creds = DummyScopedCredentials() + + with patch("google.auth.transport.requests.Request", return_value=Mock()): + credentials = _downscope_credentials(creds) + + assert credentials.scopes == ["https://www.googleapis.com/auth/sqlservice.login"] + assert credentials != creds + + def test_seconds_until_refresh_over_1_hour() -> None: """ Test _seconds_until_refresh returns proper time in seconds. diff --git a/tests/unit/test_resolver.py b/tests/unit/test_resolver.py index a7c52ed0..24e7d9ba 100644 --- a/tests/unit/test_resolver.py +++ b/tests/unit/test_resolver.py @@ -293,4 +293,41 @@ async def test_DnsResolver_global_region_skips_direct_resolution() -> None: await resolver.resolve(dns_name) # Verify mock_client was NOT called because direct resolution was skipped - mock_client.resolve_connect_settings.assert_not_called() \ No newline at end of file + mock_client.resolve_connect_settings.assert_not_called() + + +async def test_DnsResolver_no_client_error() -> None: + """Test DnsResolver throws ValueError if client is not configured for PSC DNS.""" + dns_name = "0123456789ab.fedcba9876543.europe-north2.sql-psc.goog" + resolver = DnsResolver(client=None) + with pytest.raises(ValueError) as exc_info: + await resolver.resolve(dns_name) + assert "SQLAdmin client is not configured in the resolver." in str(exc_info.value) + + +async def test_DnsResolver_invalid_domain() -> None: + """Test DnsResolver throws ValueError if input is neither connection name nor valid domain.""" + resolver = DnsResolver() + with pytest.raises(ValueError) as exc_info: + await resolver.resolve("invalidname") + assert "must have format: PROJECT:REGION:INSTANCE or be a valid DNS domain name" in str(exc_info.value) + + +async def test_DnsResolver_max_depth_reached() -> None: + """Test DnsResolver throws DnsResolutionError if max resolution depth is reached.""" + dns_name = "name0.example.com" + cname_map = {f"name{i}.example.com": f"name{i+1}.example.com" for i in range(10)} + + resolver = DnsResolver() + + async def mock_resolve_cname(name: str) -> str: + if name in cname_map: + return cname_map[name] + raise DnsResolutionError("No CNAME") + + with patch.object(resolver, "resolve_cname", AsyncMock(side_effect=mock_resolve_cname)), patch.object( + resolver, "resolve_txt", AsyncMock(side_effect=DnsResolutionError("No TXT")) + ): + with pytest.raises(DnsResolutionError) as exc_info: + await resolver.resolve(dns_name) + assert "max resolution depth reached" in str(exc_info.value) \ No newline at end of file