chore: add tests

google/cloud/alloydb/connector/client.py

@@ -70,7 +70,9 @@ def __init__(
         self._client = client if client else aiohttp.ClientSession(headers=headers)
         self._credentials = credentials
         self._alloydb_api_endpoint = alloydb_api_endpoint
-        self._user_agent = user_agent
+        # asyncpg does not currently support using metadata exchange
+        # only use metadata exchange for pg8000 driver
+        self._use_metadata = True if driver == "pg8000" else False
 
     async def _get_metadata(
         self,
@@ -150,13 +152,10 @@ async def _get_client_certificate(
 
         url = f"{self._alloydb_api_endpoint}/{API_VERSION}/projects/{project}/locations/{region}/clusters/{cluster}:generateClientCertificate"
 
-        # asyncpg does not currently support using metadata exchange
-        # only use metadata exchange for pg8000 driver
-        use_metadata = self._user_agent.endswith("pg8000")
         data = {
             "publicKey": pub_key,
             "certDuration": "3600s",
-            "useMetadataExchange": use_metadata,
+            "useMetadataExchange": self._use_metadata,
         }
 
         resp = await self._client.post(

tests/unit/test_client.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import json
-from typing import Any
+from typing import Any, Optional
 
 from aiohttp import web
 from mocks import FakeCredentials
@@ -102,3 +102,52 @@ async def test_AlloyDBClient_init_(credentials: FakeCredentials) -> None:
     assert client._client.headers["x-goog-user-project"] == "my-quota-project"
     # close client
     await client.close()
+
+
+@pytest.mark.parametrize(
+    "driver",
+    [None, "pg8000", "asyncpg"],
+)
+@pytest.mark.asyncio
+async def test_AlloyDBClient_user_agent(
+    driver: Optional[str], credentials: FakeCredentials
+) -> None:
+    """
+    Test to check whether the __init__ method of AlloyDBClient
+    properly sets user agent when passed a database driver.
+    """
+    client = AlloyDBClient(
+        "www.test-endpoint.com", "my-quota-project", credentials, driver=driver
+    )
+    if driver is None:
+        assert (
+            client._client.headers["User-Agent"]
+            == f"alloydb-python-connector/{version}"
+        )
+    else:
+        assert (
+            client._client.headers["User-Agent"]
+            == f"alloydb-python-connector/{version}+{driver}"
+        )
+    # close client
+    await client.close()
+
+
+@pytest.mark.parametrize(
+    "driver, expected",
+    [(None, False), ("pg8000", True), ("asyncpg", False)],
+)
+@pytest.mark.asyncio
+async def test_AlloyDBClient_use_metadata(
+    driver: Optional[str], expected: bool, credentials: FakeCredentials
+) -> None:
+    """
+    Test to check whether the __init__ method of AlloyDBClient
+    properly sets use_metadata.
+    """
+    client = AlloyDBClient(
+        "www.test-endpoint.com", "my-quota-project", credentials, driver=driver
+    )
+    assert client._use_metadata == expected
+    # close client
+    await client.close()