add tests

Author: sungwy

pyiceberg/catalog/rest/__init__.py

@@ -135,6 +135,7 @@ class IdentifierKind(Enum):
 OAUTH2_SERVER_URI = "oauth2-server-uri"
 SNAPSHOT_LOADING_MODE = "snapshot-loading-mode"
 AUTH = "auth"
+CUSTOM = "custom"
 
 NAMESPACE_SEPARATOR = b"\x1f".decode(UTF8)
 
@@ -249,15 +250,19 @@ def _create_session(self) -> Session:
                     session.cert = ssl_client_cert
 
         if auth_config := self.properties.get(AUTH):
-            # set up auth_manager based on the properties
             auth_type = auth_config.get("type")
             if auth_type is None:
                 raise ValueError("auth.type must be defined")
             auth_type_config = auth_config.get(auth_type, {})
-            if auth_impl := auth_config.get("impl"):
-                session.auth = AuthManagerAdapter(AuthManagerFactory.create(auth_impl, auth_type_config))
-            else:
-                session.auth = AuthManagerAdapter(AuthManagerFactory.create(auth_type, auth_type_config))
+            auth_impl = auth_config.get("impl")
+
+            if auth_type is CUSTOM and not auth_impl:
+                raise ValueError("auth.impl must be specified when using custom auth.type")
+
+            if auth_type is not CUSTOM and auth_impl:
+                raise ValueError("auth.impl can only be specified when using custom auth.type")
+
+            session.auth = AuthManagerAdapter(AuthManagerFactory.create(auth_impl or auth_type, auth_type_config))
         else:
             session.auth = AuthManagerAdapter(self._create_legacy_oauth2_auth_manager(session))
 

pyiceberg/catalog/rest/auth.py

@@ -189,8 +189,23 @@ def get_token(self) -> str:
 class OAuth2AuthManager(AuthManager):
     """Auth Manager implementation that supports OAuth2 as defined in IETF RFC6749."""
 
-    def __init__(self, token_provider: OAuth2TokenProvider):
-        self.token_provider = token_provider
+    def __init__(
+        self,
+        client_id: str,
+        client_secret: str,
+        token_url: str,
+        scope: Optional[str] = None,
+        refresh_margin: int = 60,
+        expires_in: Optional[int] = None,
+    ):
+        self.token_provider = OAuth2TokenProvider(
+            client_id,
+            client_secret,
+            token_url,
+            scope,
+            refresh_margin,
+            expires_in,
+        )
 
     def auth_header(self) -> str:
         return f"Bearer {self.token_provider.get_token()}"
@@ -274,3 +289,4 @@ def create(cls, class_or_name: str, config: Dict[str, Any]) -> AuthManager:
 AuthManagerFactory.register("noop", NoopAuthManager)
 AuthManagerFactory.register("basic", BasicAuthManager)
 AuthManagerFactory.register("legacyoauth2", LegacyOAuth2AuthManager)
+AuthManagerFactory.register("oauth2", OAuth2AuthManager)

tests/catalog/test_rest.py

@@ -1536,7 +1536,7 @@ def test_rest_catalog_with_basic_auth_type() -> None:
     assert "BasicAuthManager.__init__() missing 1 required positional argument: 'password'" in str(e.value)
 
 
-def test_rest_catalog_with_auth_impl() -> None:
+def test_rest_catalog_with_custom_auth_type() -> None:
     # Given
     catalog_properties = {
         "uri": TEST_URI,
@@ -1555,6 +1555,95 @@ def test_rest_catalog_with_auth_impl() -> None:
     assert "Could not load AuthManager class for 'dummy.nonexistent.package'" in str(e.value)
 
 
+def test_rest_catalog_with_custom_auth_type_no_impl() -> None:
+    # Given
+    catalog_properties = {
+        "uri": TEST_URI,
+        "auth": {
+            "type": "custom",
+            "custom": {
+                "property1": "one",
+                "property2": "two",
+            },
+        },
+    }
+    with pytest.raises(ValueError) as e:
+        # Missing namespace
+        RestCatalog("rest", **catalog_properties)  # type: ignore
+    assert "auth.impl must be specified when using custom auth.type" in str(e.value)
+
+
+def test_rest_catalog_with_non_custom_auth_type_impl() -> None:
+    # Given
+    catalog_properties = {
+        "uri": TEST_URI,
+        "auth": {
+            "type": "oauth2",
+            "impl": "oauth2.package",
+            "oauth2": {
+                "property1": "one",
+                "property2": "two",
+            },
+        },
+    }
+    with pytest.raises(ValueError) as e:
+        # Missing namespace
+        RestCatalog("rest", **catalog_properties)  # type: ignore
+    assert "auth.impl can only be specified when using custom auth.type" in str(e.value)
+
+
+def test_rest_catalog_with_unsupported_auth_type() -> None:
+    # Given
+    catalog_properties = {
+        "uri": TEST_URI,
+        "auth": {
+            "type": "unsupported",
+            "unsupported": {
+                "property1": "one",
+                "property2": "two",
+            },
+        },
+    }
+    with pytest.raises(ValueError) as e:
+        # Missing namespace
+        RestCatalog("rest", **catalog_properties)  # type: ignore
+    assert "Could not load AuthManager class for 'unsupported'" in str(e.value)
+
+
+def test_rest_catalog_with_oauth2_auth_type(requests_mock: Mocker) -> None:
+    requests_mock.post(
+        f"{TEST_URI}oauth2/token",
+        json={
+            "access_token": "MTQ0NjJkZmQ5OTM2NDE1ZTZjNGZmZjI3",
+            "token_type": "Bearer",
+            "expires_in": 3600,
+            "refresh_token": "IwOGYzYTlmM2YxOTQ5MGE3YmNmMDFkNTVk",
+            "scope": "read",
+        },
+        status_code=200,
+    )
+    requests_mock.get(
+        f"{TEST_URI}v1/config",
+        json={"defaults": {}, "overrides": {}},
+        status_code=200,
+    )
+    # Given
+    catalog_properties = {
+        "uri": TEST_URI,
+        "auth": {
+            "type": "oauth2",
+            "oauth2": {
+                "client_id": "some_client_id",
+                "client_secret": "some_client_secret",
+                "token_url": f"{TEST_URI}oauth2/token",
+                "scope": "read",
+            },
+        },
+    }
+    catalog = RestCatalog("rest", **catalog_properties)  # type: ignore
+    assert catalog.uri == TEST_URI
+
+
 EXAMPLE_ENV = {"PYICEBERG_CATALOG__PRODUCTION__URI": TEST_URI}