Repository: GoogleCloudPlatform/cloud-sql-python-connector
Pull request #1031 (2024-03-07 20:24:31)

refactor: update IPTypes
Aligning approach as close as possible to https://github.com/GoogleCloudPlatform/alloydb-python-connector/pull/267

Using built-in `__missing__` to better raise custom error message: https://stackoverflow.com/questions/68899057/can-one-add-a-custom-error-in-enum-to-show-valid-values


Patch:
diff --git a/.flake8 b/.flake8
index db54f31..b1d9f40 100644
--- a/.flake8
+++ b/.flake8
@@ -1,5 +1,5 @@
 [flake8]
-ignore = E203, E266, E501, W503, ANN101, ANN401
+ignore = E203, E266, E501, W503, ANN101, ANN102, ANN401
 exclude =
   # Exclude generated code.
   **/proto/**
diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py
index 760db2c..f19e299 100755
--- a/google/cloud/sql/connector/connector.py
+++ b/google/cloud/sql/connector/connector.py
@@ -48,10 +48,10 @@ ASYNC_DRIVERS = ["asyncpg"]
 class Connector:
     """A class to configure and create connections to Cloud SQL instances.
 
-    :type ip_type: IPTypes
+    :type ip_type: str | IPTypes
     :param ip_type
-        The IP type (public or private)  used to connect. IP types
-        can be either IPTypes.PUBLIC or IPTypes.PRIVATE.
+        The IP type used to connect. IP types can be either IPTypes.PUBLIC
+        ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"), or IPTypes.PSC ("PSC").
 
     :type enable_iam_auth: bool
     :param enable_iam_auth
@@ -135,7 +135,7 @@ class Connector:
         self._user_agent = user_agent
         # if ip_type is str, convert to IPTypes enum
         if isinstance(ip_type, str):
-            ip_type = IPTypes._get_ip_type_from_str(ip_type)
+            ip_type = IPTypes._from_str(ip_type)
         self._ip_type = ip_type
 
     def connect(
@@ -257,7 +257,7 @@ class Connector:
         ip_type = kwargs.pop("ip_type", self._ip_type)
         # if ip_type is str, convert to IPTypes enum
         if isinstance(ip_type, str):
-            ip_type = IPTypes._get_ip_type_from_str(ip_type)
+            ip_type = IPTypes._from_str(ip_type)
         kwargs["timeout"] = kwargs.get("timeout", self._timeout)
 
         # Host and ssl options come from the certificates and metadata, so we don't
diff --git a/google/cloud/sql/connector/instance.py b/google/cloud/sql/connector/instance.py
index 5d8f412..86e33f1 100644
--- a/google/cloud/sql/connector/instance.py
+++ b/google/cloud/sql/connector/instance.py
@@ -64,20 +64,19 @@ class IPTypes(Enum):
     PRIVATE: str = "PRIVATE"
     PSC: str = "PSC"
 
-    @staticmethod
-    def _get_ip_type_from_str(ip_type_str: str) -> IPTypes:
-        """Utility method to convert IP type from a str into IPTypes."""
-        if ip_type_str.lower() == "public":
-            ip_type = IPTypes.PUBLIC
-        elif ip_type_str.lower() == "private":
-            ip_type = IPTypes.PRIVATE
-        elif ip_type_str.lower() == "psc":
-            ip_type = IPTypes.PSC
-        else:
-            raise ValueError(
-                f"Incorrect value for ip_type, got '{ip_type_str}'. Want one of: 'public', 'private' or 'psc'."
-            )
-        return ip_type
+    @classmethod
+    def _missing_(cls, value: object) -> None:
+        raise ValueError(
+            f"Incorrect value for ip_type, got '{value}'. Want one of: "
+            f"{', '.join([repr(m.value) for m in cls])}, 'PUBLIC'."
+        )
+
+    @classmethod
+    def _from_str(cls, ip_type_str: str) -> IPTypes:
+        """Convert IP type from a str into IPTypes."""
+        if ip_type_str.upper() == "PUBLIC":
+            ip_type_str = "PRIMARY"
+        return cls(ip_type_str.upper())
 
 
 class ConnectionInfo:


Test patch:
diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py
index 08e9aa5..c16439e 100644
--- a/tests/unit/test_connector.py
+++ b/tests/unit/test_connector.py
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 import asyncio
+from typing import Union
 
 from google.auth.credentials import Credentials
 from mock import patch
@@ -131,20 +132,56 @@ async def test_Connector_Init_async_context_manager(
 
 
 @pytest.mark.parametrize(
-    "ip_type",
-    ["public", "private", "psc", "PUBLIC", "PRIVATE", "PSC"],
+    "ip_type, expected",
+    [
+        (
+            "private",
+            IPTypes.PRIVATE,
+        ),
+        (
+            "PRIVATE",
+            IPTypes.PRIVATE,
+        ),
+        (
+            IPTypes.PRIVATE,
+            IPTypes.PRIVATE,
+        ),
+        (
+            "public",
+            IPTypes.PUBLIC,
+        ),
+        (
+            "PUBLIC",
+            IPTypes.PUBLIC,
+        ),
+        (
+            IPTypes.PUBLIC,
+            IPTypes.PUBLIC,
+        ),
+        (
+            "psc",
+            IPTypes.PSC,
+        ),
+        (
+            "PSC",
+            IPTypes.PSC,
+        ),
+        (
+            IPTypes.PSC,
+            IPTypes.PSC,
+        ),
+    ],
 )
-def test_Connector_Init_ip_type_str(
-    ip_type: str, fake_credentials: Credentials
+def test_Connector_init_ip_type(
+    ip_type: Union[str, IPTypes], expected: IPTypes, fake_credentials: Credentials
 ) -> None:
-    """Test that Connector properly sets ip_type when given str."""
-    with Connector(ip_type=ip_type, credentials=fake_credentials) as connector:
-        if ip_type.lower() == "public":
-            assert connector._ip_type == IPTypes.PUBLIC
-        if ip_type.lower() == "private":
-            assert connector._ip_type == IPTypes.PRIVATE
-        if ip_type.lower() == "psc":
-            assert connector._ip_type == IPTypes.PSC
+    """
+    Test to check whether the __init__ method of Connector
+    properly sets ip_type.
+    """
+    connector = Connector(credentials=fake_credentials, ip_type=ip_type)
+    assert connector._ip_type == expected
+    connector.close()
 
 
 def test_Connector_Init_bad_ip_type(fake_credentials: Credentials) -> None:
@@ -154,7 +191,8 @@ def test_Connector_Init_bad_ip_type(fake_credentials: Credentials) -> None:
         Connector(ip_type=bad_ip_type, credentials=fake_credentials)
     assert (
         exc_info.value.args[0]
-        == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'public', 'private' or 'psc'."
+        == f"Incorrect value for ip_type, got '{bad_ip_type.upper()}'. "
+        "Want one of: 'PRIMARY', 'PRIVATE', 'PSC', 'PUBLIC'."
     )
 
 
@@ -176,7 +214,8 @@ def test_Connector_connect_bad_ip_type(
             )
         assert (
             exc_info.value.args[0]
-            == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'public', 'private' or 'psc'."
+            == f"Incorrect value for ip_type, got '{bad_ip_type.upper()}'. "
+            "Want one of: 'PRIMARY', 'PRIVATE', 'PSC', 'PUBLIC'."
         )