chore: add local dns server and tests

Author: jackwotherspoon

google/cloud/sql/connector/resolver.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dns.asyncresolver import Resolver
+import dns.asyncresolver
 
 from google.cloud.sql.connector.connection_name import _parse_instance_connection_name
 from google.cloud.sql.connector.connection_name import ConnectionName
@@ -26,7 +26,7 @@ async def resolve(self, connection_name: str) -> ConnectionName:
         return _parse_instance_connection_name(connection_name)
 
 
-class DnsResolver(Resolver):
+class DnsResolver(dns.asyncresolver.Resolver):
     """
     DnsResolver resolves domain names into instance connection names using
     TXT records in DNS.
@@ -58,7 +58,7 @@ async def query_dns(self, dns: str) -> ConnectionName:
                     continue
             # If all records failed to parse, throw error
             raise DnsResolutionError(
-                f"Unable to parse TXT record for `{dns}` -> {rdata[0]}"
+                f"Unable to parse TXT record for `{dns}` -> `{rdata[0]}`"
             )
         # Don't override above DnsResolutionError
         except DnsResolutionError:

requirements-test.txt

@@ -11,3 +11,4 @@ asyncpg==0.30.0
 python-tds==1.16.0
 aioresponses==0.7.7
 pytest-aiohttp==1.0.5
+dnserver==0.4.0

tests/conftest.py

@@ -21,6 +21,7 @@
 from typing import Any, AsyncGenerator, Generator
 
 from aiohttp import web
+from dnserver import DNSServer
 import pytest  # noqa F401 Needed to run the tests
 from unit.mocks import FakeCredentials  # type: ignore
 from unit.mocks import FakeCSQLInstance  # type: ignore
@@ -151,3 +152,13 @@ async def cache(fake_client: CloudSQLClient) -> AsyncGenerator[RefreshAheadCache
     )
     yield cache
     await cache.close()
+
+
+@pytest.fixture(autouse=True, scope="session")
+def dns_server() -> Generator:
+    """Setup local DNS server for tests with TXT records."""
+    server = DNSServer.from_toml("tests/test_zones.toml", port=5053, upstream=None)
+    server.start()
+    assert server.is_running
+    yield server
+    server.stop()

tests/test_zones.toml

@@ -0,0 +1,14 @@
+[[zones]]
+host = 'db.example.com'
+type = 'TXT'
+answer = "test-project:test-region:test-instance"
+
+[[zones]]
+host = 'db.example.com'
+type = 'TXT'
+answer = "test-project2:test-region2:test-instance2"
+
+[[zones]]
+host = 'bad.example.com'
+type = 'TXT'
+answer = "bad-instance-name"

tests/unit/test_resolver.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
+
 from google.cloud.sql.connector.connection_name import ConnectionName
+from google.cloud.sql.connector.exceptions import DnsResolutionError
 from google.cloud.sql.connector.resolver import DefaultResolver
 from google.cloud.sql.connector.resolver import DnsResolver
 
@@ -32,3 +35,38 @@ async def test_DnsResolver_with_conn_str() -> None:
     resolver = DnsResolver()
     result = await resolver.resolve(conn_str)
     assert result == conn_name
+
+
+async def test_DnsResolver_with_dns_name() -> None:
+    """Test DnsResolver resolves TXT record into proper instance connection name."""
+    resolver = DnsResolver()
+    resolver.port = 5053
+    result = await resolver.resolve(conn_str)
+    assert result == conn_name
+
+
+async def test_DnsResolver_with_malformed_txt() -> None:
+    """Test DnsResolver with TXT record that holds malformed instance connection name.
+
+    Should throw DnsResolutionError
+    """
+    resolver = DnsResolver()
+    resolver.port = 5053
+    with pytest.raises(DnsResolutionError) as exc_info:
+        await resolver.resolve("bad.example.com")
+    assert (
+        exc_info.value.args[0]
+        == "Unable to parse TXT record for `bad.example.com` -> `bad-instance-name`"
+    )
+
+
+async def test_DnsResolver_with_bad_dns_name() -> None:
+    """Test DnsResolver with bad dns name.
+
+    Should throw DnsResolutionError
+    """
+    resolver = DnsResolver()
+    resolver.port = 5053
+    with pytest.raises(DnsResolutionError) as exc_info:
+        await resolver.resolve("bad.dns.com")
+    assert exc_info.value.args[0] == "Unable to resolve TXT record for `bad.dns.com`"