Add test cases for credentials provider

test/client_test.py

@@ -8,6 +8,7 @@
 from nacos import files
 from nacos.listener import SubscribeListener, SimpleListenerManager
 from nacos.timer import NacosTimer, NacosTimerManager
+from nacos.auth import CredentialsProvider, Credentials
 import time
 import shutil
 
@@ -25,6 +26,13 @@
 # Set the following option if http requests need through by proxy
 # client.set_options(proxies={"http":"192.168.56.1:809"})
 
+class CustomCredentialsProvider(CredentialsProvider):
+    def __init__(self, ak="", sk="", token=""):
+        self.credential = Credentials(ak, sk, token)
+
+    def get_credentials(self):
+        return self.credential
+
 class TestClient(unittest.TestCase):
     def test_get_server(self):
         self.assertEqual(client.get_server(), (SERVER_1, 8848))
@@ -315,7 +323,7 @@ def test_auth_off(self):
         self.assertEqual(0, len(headers))
         self.assertEqual(0, len(params))
 
-    def test_inject_auth_info_of_config(self):
+    def test_inject_auth_info_of_config_with_fixed_ak(self):
         headers = {}
         params = {"tenant": "abc", "group": "bbb"}
         client_auth = nacos.NacosClient(SERVER_ADDRESSES, ak="1", sk="1")
@@ -325,7 +333,7 @@ def test_inject_auth_info_of_config(self):
         self.assertTrue("timeStamp" in headers)
         self.assertTrue("Spas-Signature" in headers)
 
-    def test_inject_auth_info_of_naming(self):
+    def test_inject_auth_info_of_naming_with_fixed_ak(self):
         headers = {}
         params = {"serviceName": "abc", "groupName": "bbb"}
         client_auth = nacos.NacosClient(SERVER_ADDRESSES, ak="1", sk="1")
@@ -335,6 +343,48 @@ def test_inject_auth_info_of_naming(self):
         self.assertTrue("data" in params)
         self.assertTrue("signature" in params)
 
+    def test_inject_auth_info_of_config_with_provider(self):
+        headers = {}
+        params = {"tenant": "abc", "group": "bbb"}
+
+        # access_key
+        client_auth = nacos.NacosClient(SERVER_ADDRESSES, credentials_provider=CustomCredentialsProvider("1", "1"))
+        self.assertFalse(client_auth.auth_enabled is None)
+        client_auth._inject_auth_info(headers, params, data=None, module="config")
+        self.assertEqual("1", headers.get("Spas-AccessKey"))
+        self.assertTrue("timeStamp" in headers)
+        self.assertTrue("Spas-Signature" in headers)
+        self.assertTrue("Spas-SecurityToken" not in headers)
+
+        # security_token
+        client_auth = nacos.NacosClient(SERVER_ADDRESSES, credentials_provider=CustomCredentialsProvider("1", "1", "1"))
+        self.assertFalse(client_auth.auth_enabled is None)
+        client_auth._inject_auth_info(headers, params, data=None, module="config")
+        self.assertEqual("1", headers.get("Spas-AccessKey"))
+        self.assertEqual("1", headers.get("Spas-SecurityToken"))
+        self.assertTrue("timeStamp" in headers)
+        self.assertTrue("Spas-Signature" in headers)
+
+    def test_inject_auth_info_of_naming_with_provider(self):
+        headers = {}
+        params = {"serviceName": "abc", "groupName": "bbb"}
+        # access_key
+        client_auth = nacos.NacosClient(SERVER_ADDRESSES, credentials_provider=CustomCredentialsProvider("1", "1"))
+        self.assertFalse(client_auth.auth_enabled is None)
+        client_auth._inject_auth_info(headers, params, data=None, module="naming")
+        self.assertEqual("1", params.get("ak"))
+        self.assertTrue("data" in params)
+        self.assertTrue("signature" in params)
+        self.assertTrue("Spas-SecurityToken" not in headers)
+
+        # security_token
+        client_auth = nacos.NacosClient(SERVER_ADDRESSES, credentials_provider=CustomCredentialsProvider("1", "1", "1"))
+        self.assertFalse(client_auth.auth_enabled is None)
+        client_auth._inject_auth_info(headers, params, data=None, module="naming")
+        self.assertEqual("1", params.get("ak"))
+        self.assertEqual("1", params.get("Spas-SecurityToken"))
+        self.assertTrue("data" in params)
+        self.assertTrue("signature" in params)
 
 if __name__ == '__main__':
     unittest.main()

test/client_v2_test.py

@@ -8,6 +8,7 @@
 from v2.nacos.naming.model.naming_param import RegisterInstanceParam, DeregisterInstanceParam, \
     BatchRegisterInstanceParam, GetServiceParam, ListServiceParam, SubscribeServiceParam, ListInstanceParam
 from v2.nacos.naming.nacos_naming_service import NacosNamingService
+from v2.nacos.common.auth import CredentialsProvider, Credentials
 
 client_config = (ClientConfigBuilder()
                  .access_key(os.getenv('NACOS_ACCESS_KEY'))
@@ -17,10 +18,16 @@
                  .grpc_config(GRPCConfig(grpc_timeout=5000))
                  .build())
 
+class CustomCredentialsProvider(CredentialsProvider):
+    def __init__(self, ak="", sk="", token=""):
+        self.credential = Credentials(ak, sk, token)
+
+    def get_credentials(self):
+        return self.credential
 
 class TestClientV2(unittest.IsolatedAsyncioTestCase):
 
-    async def test_register_with_endpoint(self):
+    async def test_register_with_endpoint_and_fixed_ak(self):
         config = (ClientConfigBuilder()
                          .access_key(os.getenv('NACOS_ACCESS_KEY'))
                          .secret_key(os.getenv('NACOS_SECRET_KEY'))
@@ -32,6 +39,17 @@ async def test_register_with_endpoint(self):
         client = await NacosNamingService.create_naming_service(config)
         assert await client.server_health()
 
+    async def test_register_with_endpoint_and_provider(self):
+        config = (ClientConfigBuilder()
+                         .credentials_provider(CustomCredentialsProvider(os.getenv('NACOS_ACCESS_KEY'), os.getenv('NACOS_SECRET_KEY')))
+                         .endpoint(os.getenv('NACOS_SERVER_ENDPOINT', 'localhost:8848'))
+                         .log_level('INFO')
+                         .grpc_config(GRPCConfig(grpc_timeout=5000))
+                         .build())
+
+        client = await NacosNamingService.create_naming_service(config)
+        assert await client.server_health()
+
     async def test_register(self):
         client = await NacosNamingService.create_naming_service(client_config)
         assert await client.server_health()

test/config_client_v2_test.py

@@ -6,6 +6,7 @@
 from v2.nacos.common.client_config_builder import ClientConfigBuilder
 from v2.nacos.config.model.config_param import ConfigParam
 from v2.nacos.config.nacos_config_service import NacosConfigService
+from v2.nacos.common.auth import CredentialsProvider, Credentials
 
 client_config = (ClientConfigBuilder()
                  .access_key(os.getenv('NACOS_ACCESS_KEY'))
@@ -15,6 +16,12 @@
                  .grpc_config(GRPCConfig(grpc_timeout=5000))
                  .build())
 
+class CustomCredentialsProvider(CredentialsProvider):
+    def __init__(self, ak="", sk="", token=""):
+        self.credential = Credentials(ak, sk, token)
+
+    def get_credentials(self):
+        return self.credential
 
 class TestClientV2(unittest.IsolatedAsyncioTestCase):
     async def test_publish_config(self):
@@ -235,15 +242,7 @@ async def config_listener(tenant, data_id, group, content):
 
         await client.shutdown()
 
-    async def test_gray_config(self):
-        client_cfg = (ClientConfigBuilder()
-                      .access_key(os.getenv('NACOS_ACCESS_KEY'))
-                      .secret_key(os.getenv('NACOS_SECRET_KEY'))
-                      .server_address(os.getenv('NACOS_SERVER_ADDR', 'localhost:8848'))
-                      .log_level('INFO')
-                      .app_conn_labels({"k1": "v1", "k2": "v2", "nacos_config_gray_label": "gray"})
-                      .grpc_config(GRPCConfig(grpc_timeout=5000))
-                      .build())
+    async def _gray_config(self, client_cfg):
         client = await NacosConfigService.create_config_service(client_cfg)
 
         dataID = "com.alibaba.nacos.test.config.gray"
@@ -255,3 +254,24 @@ async def config_listener(tenant, data_id, group, content):
         await client.add_listener(dataID, groupName, config_listener)
 
         await asyncio.sleep(1000)
+
+    async def test_gray_config_with_fixed_ak(self):
+        client_cfg = (ClientConfigBuilder()
+                      .access_key(os.getenv('NACOS_ACCESS_KEY'))
+                      .secret_key(os.getenv('NACOS_SECRET_KEY'))
+                      .server_address(os.getenv('NACOS_SERVER_ADDR', 'localhost:8848'))
+                      .log_level('INFO')
+                      .app_conn_labels({"k1": "v1", "k2": "v2", "nacos_config_gray_label": "gray"})
+                      .grpc_config(GRPCConfig(grpc_timeout=5000))
+                      .build())
+        await self._gray_config(client_cfg)
+
+    async def test_gray_config_with_provider(self):
+        client_cfg = (ClientConfigBuilder()
+                      .credentials_provider(CustomCredentialsProvider(os.getenv('NACOS_ACCESS_KEY'), os.getenv('NACOS_SECRET_KEY')))
+                      .server_address(os.getenv('NACOS_SERVER_ADDR', 'localhost:8848'))
+                      .log_level('INFO')
+                      .app_conn_labels({"k1": "v1", "k2": "v2", "nacos_config_gray_label": "gray"})
+                      .grpc_config(GRPCConfig(grpc_timeout=5000))
+                      .build())
+        await self._gray_config(client_cfg)