[#4208] improvement(client-python): Add Integration Test of OAuth2Tok…

…enProvider of client-python (#4663)

### What changes were proposed in this pull request?

* Add integration test of `OAuth2TokenProvider`, following the steps
mentioned in OAuth section of
[`/docs/security/how-to-authenticate.md`](https://github.com/apache/gravitino/blob/main/docs/security/how-to-authenticate.md)
* Add some base class for containers and tests for future development
* Remove some duplicate code

### Why are the changes needed?

Fix: #4208 

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

ITs added, tests with `./gradlew client:client-python:test
-PskipDockerTests=false`

---------

Co-authored-by: TimWang <[email protected]>

clients/client-python/build.gradle.kts

@@ -223,6 +223,7 @@ tasks {
       "START_EXTERNAL_GRAVITINO" to "true",
       "DOCKER_TEST" to dockerTest.toString(),
       "GRAVITINO_CI_HIVE_DOCKER_IMAGE" to "apache/gravitino-ci:hive-0.1.13",
+      "GRAVITINO_OAUTH2_SAMPLE_SERVER" to "datastrato/sample-authorization-server:0.3.0",
       // Set the PYTHONPATH to the client-python directory, make sure the tests can import the
       // modules from the client-python directory.
       "PYTHONPATH" to "${project.rootDir.path}/clients/client-python"

clients/client-python/gravitino/dto/responses/oauth2_token_response.py

@@ -44,8 +44,6 @@ def validate(self):
         Raise:
             IllegalArgumentException If the response is invalid, this exception is thrown.
         """
-        super().validate()
-
         if self._access_token is None:
             raise IllegalArgumentException("Invalid access token: None")
 

clients/client-python/gravitino/utils/http_client.py

@@ -140,7 +140,7 @@ def _make_request(self, opener, request, timeout=None) -> Tuple[bool, Response]:
         except HTTPError as err:
             err_body = err.read()
 
-            if err_body is None:
+            if err_body is None or len(err_body) == 0:
                 return (
                     False,
                     ErrorResponse.generate_error_response(RESTException, err.reason),

clients/client-python/requirements-dev.txt

@@ -29,3 +29,4 @@ cachetools==5.3.3
 readerwriterlock==1.0.9
 docker==7.1.0
 pyjwt[crypto]==2.8.0
+jwcrypto==1.5.6

clients/client-python/tests/integration/auth/__init__.py

@@ -0,0 +1,18 @@
+"""
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements.  See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership.  The ASF licenses this file
+to you under the Apache License, Version 2.0 (the
+"License"); you may not use this file except in compliance
+with the License.  You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing,
+software distributed under the License is distributed on an
+"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, either express or implied.  See the License for the
+specific language governing permissions and limitations
+under the License.
+"""

clients/client-python/tests/integration/test_simple_auth_client.py → clients/client-python/tests/integration/auth/test_auth_common.py

@@ -29,15 +29,17 @@
     Catalog,
     Fileset,
 )
-from gravitino.auth.simple_auth_provider import SimpleAuthProvider
 from gravitino.exceptions.base import GravitinoRuntimeException
-from tests.integration.integration_test_env import IntegrationTestEnv
 
 logger = logging.getLogger(__name__)
 
 
-class TestSimpleAuthClient(IntegrationTestEnv):
-    creator: str = "test_client"
+class TestCommonAuth:
+    """
+    A common test set for AuthProvider Integration Tests
+    """
+
+    creator: str = "test"
     metalake_name: str = "TestClient_metalake" + str(randint(1, 10000))
     catalog_name: str = "fileset_catalog"
     catalog_location_prop: str = "location"  # Fileset Catalog must set `location`
@@ -63,16 +65,8 @@ class TestSimpleAuthClient(IntegrationTestEnv):
         metalake_name, catalog_name, schema_name
     )
     fileset_ident: NameIdentifier = NameIdentifier.of(schema_name, fileset_name)
-
-    def setUp(self):
-        os.environ["GRAVITINO_USER"] = self.creator
-        self.gravitino_admin_client = GravitinoAdminClient(
-            uri="http://localhost:8090", auth_data_provider=SimpleAuthProvider()
-        )
-        self.init_test_env()
-
-    def tearDown(self):
-        self.clean_test_data()
+    gravitino_admin_client: GravitinoAdminClient
+    gravitino_client: GravitinoClient
 
     def clean_test_data(self):
         catalog = self.gravitino_client.load_catalog(name=self.catalog_name)
@@ -117,14 +111,7 @@ def clean_test_data(self):
         os.environ["GRAVITINO_USER"] = ""
 
     def init_test_env(self):
-        self.gravitino_admin_client.create_metalake(
-            self.metalake_name, comment="", properties={}
-        )
-        self.gravitino_client = GravitinoClient(
-            uri="http://localhost:8090",
-            metalake_name=self.metalake_name,
-            auth_data_provider=SimpleAuthProvider(),
-        )
+
         catalog = self.gravitino_client.create_catalog(
             name=self.catalog_name,
             catalog_type=Catalog.Type.FILESET,

clients/client-python/tests/integration/auth/test_oauth2_client.py

@@ -0,0 +1,178 @@
+"""
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements.  See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership.  The ASF licenses this file
+to you under the Apache License, Version 2.0 (the
+"License"); you may not use this file except in compliance
+with the License.  You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing,
+software distributed under the License is distributed on an
+"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, either express or implied.  See the License for the
+specific language governing permissions and limitations
+under the License.
+"""
+
+import os
+import subprocess
+import logging
+import unittest
+import sys
+import requests
+from jwcrypto import jwk
+
+from gravitino.auth.auth_constants import AuthConstants
+from gravitino.auth.default_oauth2_token_provider import DefaultOAuth2TokenProvider
+from gravitino import GravitinoAdminClient, GravitinoClient
+from gravitino.exceptions.base import GravitinoRuntimeException
+
+from tests.integration.auth.test_auth_common import TestCommonAuth
+from tests.integration.integration_test_env import (
+    IntegrationTestEnv,
+    check_gravitino_server_status,
+)
+from tests.integration.containers.oauth2_container import OAuth2Container
+
+logger = logging.getLogger(__name__)
+
+DOCKER_TEST = os.environ.get("DOCKER_TEST")
+
+
+@unittest.skipIf(
+    DOCKER_TEST == "false",
+    "Skipping tests when DOCKER_TEST=false",
+)
+class TestOAuth2(IntegrationTestEnv, TestCommonAuth):
+
+    oauth2_container: OAuth2Container = None
+
+    @classmethod
+    def setUpClass(cls):
+
+        cls._get_gravitino_home()
+
+        cls.oauth2_container = OAuth2Container()
+        cls.oauth2_container_ip = cls.oauth2_container.get_ip()
+
+        cls.oauth2_server_uri = f"http://{cls.oauth2_container_ip}:8177"
+
+        # Get PEM from OAuth Server
+        default_sign_key = cls._get_default_sign_key()
+
+        cls.config = {
+            "gravitino.authenticators": "oauth",
+            "gravitino.authenticator.oauth.serviceAudience": "test",
+            "gravitino.authenticator.oauth.defaultSignKey": default_sign_key,
+            "gravitino.authenticator.oauth.serverUri": cls.oauth2_server_uri,
+            "gravitino.authenticator.oauth.tokenPath": "/oauth2/token",
+        }
+
+        cls.oauth2_conf_path = f"{cls.gravitino_home}/conf/gravitino.conf"
+
+        # append the hadoop conf to server
+        cls._append_conf(cls.config, cls.oauth2_conf_path)
+        # restart the server
+        cls._restart_server_with_oauth()
+
+    @classmethod
+    def tearDownClass(cls):
+        try:
+            # reset server conf
+            cls._reset_conf(cls.config, cls.oauth2_conf_path)
+            # restart server
+            cls.restart_server()
+        finally:
+            # close oauth2 container
+            cls.oauth2_container.close()
+
+    @classmethod
+    def _get_default_sign_key(cls) -> str:
+
+        jwk_uri = f"{cls.oauth2_server_uri}/oauth2/jwks"
+
+        # Get JWK from OAuth2 Server
+        res = requests.get(jwk_uri).json()
+        key = res["keys"][0]
+
+        # Convert JWK to PEM
+        pem = jwk.JWK(**key).export_to_pem().decode("utf-8")
+
+        default_sign_key = "".join(pem.split("\n")[1:-2])
+
+        return default_sign_key
+
+    @classmethod
+    def _restart_server_with_oauth(cls):
+        logger.info("Restarting Gravitino server...")
+        gravitino_home = os.environ.get("GRAVITINO_HOME")
+        gravitino_startup_script = os.path.join(gravitino_home, "bin/gravitino.sh")
+        if not os.path.exists(gravitino_startup_script):
+            raise GravitinoRuntimeException(
+                f"Can't find Gravitino startup script: {gravitino_startup_script}, "
+                "Please execute `./gradlew compileDistribution -x test` in the Gravitino "
+                "project root directory."
+            )
+
+        # Restart Gravitino Server
+        env_vars = os.environ.copy()
+        result = subprocess.run(
+            [gravitino_startup_script, "restart"],
+            env=env_vars,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if result.stdout:
+            logger.info("stdout: %s", result.stdout)
+        if result.stderr:
+            logger.info("stderr: %s", result.stderr)
+
+        oauth2_token_provider = DefaultOAuth2TokenProvider(
+            f"{cls.oauth2_server_uri}", "test:test", "test", "oauth2/token"
+        )
+
+        auth_header = {
+            AuthConstants.HTTP_HEADER_AUTHORIZATION: oauth2_token_provider.get_token_data().decode(
+                "utf-8"
+            )
+        }
+
+        if not check_gravitino_server_status(headers=auth_header):
+            logger.error("ERROR: Can't start Gravitino server!")
+            sys.exit(0)
+
+    def setUp(self):
+        oauth2_token_provider = DefaultOAuth2TokenProvider(
+            f"{self.oauth2_server_uri}", "test:test", "test", "oauth2/token"
+        )
+
+        self.gravitino_admin_client = GravitinoAdminClient(
+            uri="http://localhost:8090", auth_data_provider=oauth2_token_provider
+        )
+
+        self.init_test_env()
+
+    def init_test_env(self):
+
+        self.gravitino_admin_client.create_metalake(
+            self.metalake_name, comment="", properties={}
+        )
+
+        oauth2_token_provider = DefaultOAuth2TokenProvider(
+            f"{self.oauth2_server_uri}", "test:test", "test", "oauth2/token"
+        )
+
+        self.gravitino_client = GravitinoClient(
+            uri="http://localhost:8090",
+            metalake_name=self.metalake_name,
+            auth_data_provider=oauth2_token_provider,
+        )
+
+        super().init_test_env()
+
+    def tearDown(self):
+        self.clean_test_data()

clients/client-python/tests/integration/auth/test_simple_auth_client.py

@@ -0,0 +1,58 @@
+"""
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements.  See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership.  The ASF licenses this file
+to you under the Apache License, Version 2.0 (the
+"License"); you may not use this file except in compliance
+with the License.  You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing,
+software distributed under the License is distributed on an
+"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, either express or implied.  See the License for the
+specific language governing permissions and limitations
+under the License.
+"""
+
+import logging
+import os
+
+from gravitino import (
+    GravitinoClient,
+    GravitinoAdminClient,
+)
+from gravitino.auth.simple_auth_provider import SimpleAuthProvider
+
+from tests.integration.auth.test_auth_common import TestCommonAuth
+from tests.integration.integration_test_env import IntegrationTestEnv
+
+logger = logging.getLogger(__name__)
+
+
+class TestSimpleAuthClient(IntegrationTestEnv, TestCommonAuth):
+
+    def setUp(self):
+        os.environ["GRAVITINO_USER"] = self.creator
+        self.gravitino_admin_client = GravitinoAdminClient(
+            uri="http://localhost:8090", auth_data_provider=SimpleAuthProvider()
+        )
+
+        self.init_test_env()
+
+    def init_test_env(self):
+        self.gravitino_admin_client.create_metalake(
+            self.metalake_name, comment="", properties={}
+        )
+        self.gravitino_client = GravitinoClient(
+            uri="http://localhost:8090",
+            metalake_name=self.metalake_name,
+            auth_data_provider=SimpleAuthProvider(),
+        )
+
+        super().init_test_env()
+
+    def tearDown(self):
+        self.clean_test_data()