Merge pull request #72 from Serverless-Devs/add-ots-kb

refactor(knowledgebase): 添加STS token支持并优化代码格式

Author: OhYee

agentrun/knowledgebase/api/__data_async_template.py

@@ -595,6 +595,7 @@ def _build_agent_storage_client(
         return AgentStorageClient(
             access_key_id=cfg.get_access_key_id(),
             access_key_secret=cfg.get_access_key_secret(),
+            sts_token=cfg.get_security_token(),
             ots_endpoint=ots_endpoint,
             ots_instance_name=self.provider_settings.ots_instance_name,
         )

agentrun/knowledgebase/api/data.py

@@ -854,6 +854,7 @@ def _build_agent_storage_client(
         return AgentStorageClient(
             access_key_id=cfg.get_access_key_id(),
             access_key_secret=cfg.get_access_key_secret(),
+            sts_token=cfg.get_security_token(),
             ots_endpoint=ots_endpoint,
             ots_instance_name=self.provider_settings.ots_instance_name,
         )

tests/unittests/knowledgebase/test_ots_knowledgebase.py

@@ -600,6 +600,7 @@ def test_build_client(self, mock_client_class):
         mock_config.get_region_id.return_value = "cn-hangzhou"
         mock_config.get_access_key_id.return_value = "test-ak"
         mock_config.get_access_key_secret.return_value = "test-sk"
+        mock_config.get_security_token.return_value = "test-sts"
 
         with patch.object(Config, "with_configs", return_value=mock_config):
             api = OTSDataAPI(
@@ -613,6 +614,7 @@ def test_build_client(self, mock_client_class):
         mock_client_class.assert_called_once_with(
             access_key_id="test-ak",
             access_key_secret="test-sk",
+            sts_token="test-sts",
             ots_endpoint="http://ots-cn-hangzhou.aliyuncs.com",
             ots_instance_name="test-instance",
         )