diff --git a/agentkit/sdk/identity/auth.py b/agentkit/sdk/identity/auth.py index 5971c08..cbfbad5 100644 --- a/agentkit/sdk/identity/auth.py +++ b/agentkit/sdk/identity/auth.py @@ -41,7 +41,7 @@ def _get_api_key() -> str: endpoint = resolve_endpoint("identity") access_key = creds.access_key secret_key = creds.secret_key - session_token = "" + session_token = getattr(creds, "session_token", None) or None response = ve_request( request_body={ @@ -49,9 +49,9 @@ def _get_api_key() -> str: "IdentityToken": "identity_token", }, action="GetResourceApiKey", - header={"X-Security-Token": session_token} if session_token else {}, ak=access_key, sk=secret_key, + session_token=session_token, version=endpoint.api_version, service=endpoint.service, host=endpoint.host, diff --git a/agentkit/toolkit/volcengine/code_pipeline.py b/agentkit/toolkit/volcengine/code_pipeline.py index c81ed8f..3fa49d6 100644 --- a/agentkit/toolkit/volcengine/code_pipeline.py +++ b/agentkit/toolkit/volcengine/code_pipeline.py @@ -31,11 +31,13 @@ def __init__( secret_key: str = "", region: str = "", provider: str | None = None, + session_token: str = "", ) -> None: # Use provided region or None to trigger auto-detection in VolcConfiguration config = VolcConfiguration( access_key=access_key or None, secret_key=secret_key or None, + session_token=session_token or None, region=region or None, provider=provider or None, ) @@ -45,6 +47,7 @@ def __init__( creds = config.get_service_credentials("cp") self.volcengine_access_key = creds.access_key self.volcengine_secret_key = creds.secret_key + self.volcengine_session_token = creds.session_token or None elif not all([access_key, secret_key]): raise ValueError( "Error create cp instance: missing access key or secret key", @@ -52,6 +55,7 @@ def __init__( else: self.volcengine_access_key = access_key self.volcengine_secret_key = secret_key + self.volcengine_session_token = session_token or None endpoint = config.get_service_endpoint("cp") @@ -61,12 +65,10 @@ def __init__( self.host = endpoint.host self.content_type = "application/json" - def _get_default_workspace(self) -> str: - logger.info("Getting default workspace...") - - res = ve_request( - request_body={}, - action="GetDefaultWorkspaceInner", + def _ve_request(self, request_body: dict, action: str) -> dict: + return ve_request( + request_body=request_body, + action=action, ak=self.volcengine_access_key, sk=self.volcengine_secret_key, service=self.service, @@ -74,6 +76,15 @@ def _get_default_workspace(self) -> str: region=self.region, host=self.host, content_type=self.content_type, + session_token=self.volcengine_session_token, + ) + + def _get_default_workspace(self) -> str: + logger.info("Getting default workspace...") + + res = self._ve_request( + request_body={}, + action="GetDefaultWorkspaceInner", ) try: @@ -134,16 +145,9 @@ def create_workspace( if visible_users: request_body["VisibleUsers"] = visible_users - res = ve_request( + res = self._ve_request( request_body=request_body, action="CreateWorkspace", - ak=self.volcengine_access_key, - sk=self.volcengine_secret_key, - service=self.service, - version=self.version, - region=self.region, - host=self.host, - content_type=self.content_type, ) try: @@ -206,16 +210,9 @@ def list_workspaces( if workspace_ids: request_body["Filter"]["Ids"] = workspace_ids - res = ve_request( + res = self._ve_request( request_body=request_body, action="ListWorkspaces", - ak=self.volcengine_access_key, - sk=self.volcengine_secret_key, - service=self.service, - version=self.version, - region=self.region, - host=self.host, - content_type=self.content_type, ) try: @@ -301,7 +298,7 @@ def _create_pipeline( parameters: list[dict[str, str]] | None = None, ) -> str: logger.info("Creating pipeline...") - res = ve_request( + res = self._ve_request( request_body={ "WorkspaceId": workspace_id, "Name": pipeline_name, @@ -309,13 +306,6 @@ def _create_pipeline( "Parameters": parameters or [], }, action="CreatePipeline", - ak=self.volcengine_access_key, - sk=self.volcengine_secret_key, - service=self.service, - version=self.version, - region=self.region, - host=self.host, - content_type=self.content_type, ) try: @@ -366,16 +356,9 @@ def run_pipeline( if resources: request_body["Resources"] = resources - res = ve_request( + res = self._ve_request( request_body=request_body, action="RunPipeline", - ak=self.volcengine_access_key, - sk=self.volcengine_secret_key, - service=self.service, - version=self.version, - region=self.region, - host=self.host, - content_type=self.content_type, ) try: @@ -459,16 +442,9 @@ def list_pipeline_runs( if run_ids: request_body["Filter"]["Ids"] = run_ids - res = ve_request( + res = self._ve_request( request_body=request_body, action="ListPipelineRuns", - ak=self.volcengine_access_key, - sk=self.volcengine_secret_key, - service=self.service, - version=self.version, - region=self.region, - host=self.host, - content_type=self.content_type, ) try: @@ -569,16 +545,9 @@ def list_pipelines( if pipeline_ids: request_body["Filter"]["Ids"] = pipeline_ids - res = ve_request( + res = self._ve_request( request_body=request_body, action="ListPipelines", - ak=self.volcengine_access_key, - sk=self.volcengine_secret_key, - service=self.service, - version=self.version, - region=self.region, - host=self.host, - content_type=self.content_type, ) try: @@ -684,16 +653,9 @@ def list_pipeline_run_stages_inner( "PipelineRunId": pipeline_run_id, } - res = ve_request( + res = self._ve_request( request_body=request_body, action="ListPipelineRunStagesInner", - ak=self.volcengine_access_key, - sk=self.volcengine_secret_key, - service=self.service, - version=self.version, - region=self.region, - host=self.host, - content_type=self.content_type, ) try: @@ -757,16 +719,9 @@ def get_task_run_log_download_uri( "StepName": step_name, } - res = ve_request( + res = self._ve_request( request_body=request_body, action="GetTaskRunLogDownloadURI", - ak=self.volcengine_access_key, - sk=self.volcengine_secret_key, - service=self.service, - version=self.version, - region=self.region, - host=self.host, - content_type=self.content_type, ) try: diff --git a/agentkit/toolkit/volcengine/cr.py b/agentkit/toolkit/volcengine/cr.py index 1e63ff4..70ebbb7 100644 --- a/agentkit/toolkit/volcengine/cr.py +++ b/agentkit/toolkit/volcengine/cr.py @@ -32,9 +32,11 @@ def __init__( secret_key: str, region: str | None = None, provider: str | None = None, + session_token: str | None = None, ): self.ak = access_key self.sk = secret_key + self.session_token = session_token or None config = VolcConfiguration(region=region or None, provider=provider or None) ep = resolve_endpoint( @@ -47,6 +49,20 @@ def __init__( self.host = ep.host self.scheme = ep.scheme + def _ve_request(self, request_body: dict, action: str) -> dict: + return ve_request( + request_body=request_body, + action=action, + ak=self.ak, + sk=self.sk, + service="cr", + version=self.version, + region=self.region, + host=self.host, + scheme=self.scheme, + session_token=self.session_token, + ) + def _create_instance( self, instance_name: str = DEFAULT_CR_INSTANCE_NAME, @@ -74,7 +90,7 @@ def _create_instance( if status != "NONEXIST": logger.debug(f"cr instance {instance_name} already running") return instance_name - response = ve_request( + response = self._ve_request( request_body={ "Name": instance_name, "ResourceTags": [ @@ -83,13 +99,6 @@ def _create_instance( "Type": instance_type, }, action="CreateRegistry", - ak=self.ak, - sk=self.sk, - service="cr", - version=self.version, - region=self.region, - host=self.host, - scheme=self.scheme, ) logger.debug(f"create cr instance {instance_name}: {response}") @@ -129,20 +138,13 @@ def _check_instance(self, instance_name: str) -> str: Returns: cr instance status """ - response = ve_request( + response = self._ve_request( request_body={ "Filter": { "Names": [instance_name], } }, action="ListRegistries", - ak=self.ak, - sk=self.sk, - service="cr", - version=self.version, - region=self.region, - host=self.host, - scheme=self.scheme, ) logger.debug(f"check cr instance {instance_name}: {response}") @@ -168,19 +170,12 @@ def _create_namespace( Returns: cr namespace name """ - response = ve_request( + response = self._ve_request( request_body={ "Name": namespace_name, "Registry": instance_name, }, action="CreateNamespace", - ak=self.ak, - sk=self.sk, - service="cr", - version=self.version, - region=self.region, - host=self.host, - scheme=self.scheme, ) logger.debug(f"create cr namespace {namespace_name}: {response}") @@ -217,7 +212,7 @@ def _create_repo( Returns: cr repo name """ - response = ve_request( + response = self._ve_request( request_body={ "Name": repo_name, "Registry": instance_name, @@ -226,13 +221,6 @@ def _create_repo( "Description": "veadk cr repo", }, action="CreateRepository", - ak=self.ak, - sk=self.sk, - service="cr", - version=self.version, - region=self.region, - host=self.host, - scheme=self.scheme, ) logger.debug(f"create cr repo {repo_name}: {response}") @@ -255,18 +243,11 @@ def _get_authorization_token(self, instance_name: str): """ get cr authorization token """ - response = ve_request( + response = self._ve_request( request_body={ "Registry": instance_name, }, action="GetAuthorizationToken", - ak=self.ak, - sk=self.sk, - service="cr", - version=self.version, - region=self.region, - host=self.host, - scheme=self.scheme, ) logger.debug("got cr authorization token") @@ -291,18 +272,11 @@ def _get_public_endpoint(self, instance_name: str = DEFAULT_CR_INSTANCE_NAME): """ get cr public endpoint """ - response = ve_request( + response = self._ve_request( request_body={ "Registry": instance_name, }, action="GetPublicEndpoint", - ak=self.ak, - sk=self.sk, - service="cr", - version=self.version, - region=self.region, - host=self.host, - scheme=self.scheme, ) logger.debug(f"get cr public endpoint: {response}") if "Error" in response["ResponseMetadata"]: @@ -318,19 +292,12 @@ def _update_public_endpoint(self, instance_name: str, enabled: bool): """ update cr public endpoint """ - response = ve_request( + response = self._ve_request( request_body={ "Registry": instance_name, "Enabled": enabled, }, action="UpdatePublicEndpoint", - ak=self.ak, - sk=self.sk, - service="cr", - version=self.version, - region=self.region, - host=self.host, - scheme=self.scheme, ) logger.debug(f"update cr public endpoint: {response}") if "Error" in response["ResponseMetadata"]: @@ -354,7 +321,7 @@ def _create_endpoint_acl_policies( """ create endpoint acl policies """ - response = ve_request( + response = self._ve_request( request_body={ "Registry": instance_name, "Type": policy_type, @@ -362,13 +329,6 @@ def _create_endpoint_acl_policies( "Description": description, }, action="CreateEndpointAclPolicies", - ak=self.ak, - sk=self.sk, - service="cr", - version=self.version, - region=self.region, - host=self.host, - scheme=self.scheme, ) logger.debug(f"create endpoint acl policies: {response}") @@ -387,18 +347,11 @@ def _list_domains(self, instance_name: str = DEFAULT_CR_INSTANCE_NAME): """ list cr domains """ - response = ve_request( + response = self._ve_request( request_body={ "Registry": instance_name, }, action="ListDomains", - ak=self.ak, - sk=self.sk, - service="cr", - version=self.version, - region=self.region, - host=self.host, - scheme=self.scheme, ) logger.debug(f"list cr domains: {response}") if "Error" in response["ResponseMetadata"]: diff --git a/agentkit/toolkit/volcengine/services/cr_service.py b/agentkit/toolkit/volcengine/services/cr_service.py index 1c4bbf1..6de6056 100644 --- a/agentkit/toolkit/volcengine/services/cr_service.py +++ b/agentkit/toolkit/volcengine/services/cr_service.py @@ -207,6 +207,7 @@ def _init_client(self, region: Optional[str] = None) -> None: self._vecr_client = ve_cr.VeCR( access_key=creds.access_key, secret_key=creds.secret_key, + session_token=creds.session_token, region=endpoint.region, provider=self.provider, ) diff --git a/agentkit/utils/ve_sign.py b/agentkit/utils/ve_sign.py index 25947ee..19c819f 100644 --- a/agentkit/utils/ve_sign.py +++ b/agentkit/utils/ve_sign.py @@ -212,19 +212,22 @@ def norm_query(params): return query.replace("+", "%20") -# 第一步：准备辅助函数。 -# sha256 非对称加密 def hmac_sha256(key: bytes, content: str): return hmac.new(key, content.encode("utf-8"), hashlib.sha256).digest() -# sha256 hash算法 def hash_sha256(content: str): return hashlib.sha256(content.encode("utf-8")).hexdigest() -# 第二步：签名请求函数 -def request(method, date, query, header, ak, sk, action, body): +def _header_value(headers: dict, name: str) -> str | None: + for key, value in headers.items(): + if key.lower() == name.lower() and value: + return str(value) + return None + + +def request(method, date, query, header, ak, sk, action, body, session_token=None): # 第三步：创建身份证明。其中的 Service 和 Region 字段是固定的。ak 和 sk 分别代表 # AccessKeyID 和 SecretAccessKey。同时需要初始化签名结构体。一些签名计算时需要的属性也在这里处理。 # 初始化身份证明结构体 @@ -257,24 +260,28 @@ def request(method, date, query, header, ak, sk, action, body): "X-Date": x_date, "Content-Type": request_param["content_type"], } + token = session_token or _header_value(header, "X-Security-Token") + if token: + sign_result["X-Security-Token"] = token + # 第五步：计算 Signature 签名。 - signed_headers_str = ";".join( - ["content-type", "host", "x-content-sha256", "x-date"] - ) - # signed_headers_str = signed_headers_str + ";x-security-token" + signed_header_names = ["content-type", "host", "x-content-sha256", "x-date"] + canonical_header_lines = [ + "content-type:" + request_param["content_type"], + "host:" + request_param["host"], + "x-content-sha256:" + x_content_sha256, + "x-date:" + x_date, + ] + if token: + signed_header_names.append("x-security-token") + canonical_header_lines.append("x-security-token:" + token) + signed_headers_str = ";".join(signed_header_names) canonical_request_str = "\n".join( [ request_param["method"].upper(), request_param["path"], norm_query(request_param["query"]), - "\n".join( - [ - "content-type:" + request_param["content_type"], - "host:" + request_param["host"], - "x-content-sha256:" + x_content_sha256, - "x-date:" + x_date, - ] - ), + "\n".join(canonical_header_lines), "", signed_headers_str, x_content_sha256, @@ -309,9 +316,14 @@ def request(method, date, query, header, ak, sk, action, body): signature, ) ) + if token: + header = { + key: value + for key, value in header.items() + if key.lower() != "x-security-token" + } header = ensure_x_custom_source_header(header) header = {**header, **sign_result} - # header = {**header, **{"X-Security-Token": SessionToken}} # 第六步：将 Signature 签名写入 HTTP Header 中，并发送 HTTP 请求。 r = _signed_request( method=method, @@ -335,6 +347,7 @@ def ve_request( header: dict | None = None, content_type: str = "application/json", scheme: str = "https", + session_token: str | None = None, ): # response_body = request("Get", datetime.datetime.utcnow(), {}, {}, AK, SK, "ListUsers", None) # print(response_body) @@ -355,14 +368,22 @@ def ve_request( AK = ak SK = sk - now = datetime.datetime.utcnow() + now = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) # Body的格式需要配合Content-Type，API使用的类型请阅读具体的官方文档，如:json格式需要json.dumps(obj) # response_body = request("GET", now, {"Limit": "2"}, {}, AK, SK, "ListUsers", None) import json response_body = request( - "POST", now, {}, header or {}, AK, SK, action, json.dumps(request_body) + "POST", + now, + {}, + header or {}, + AK, + SK, + action, + json.dumps(request_body), + session_token=session_token, ) check_error(response_body) return response_body diff --git a/tests/auth/test_identity_auth.py b/tests/auth/test_identity_auth.py new file mode 100644 index 0000000..ec4a36d --- /dev/null +++ b/tests/auth/test_identity_auth.py @@ -0,0 +1,102 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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. + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + +import pytest + +import agentkit.sdk.identity.auth as identity_auth +from agentkit.sdk.identity.auth import requires_api_key +from agentkit.toolkit.errors import ApiError + + +@dataclass +class _Creds: + access_key: str = "AK" + secret_key: str = "SK" + session_token: str | None = "STS_TOKEN" + + +@dataclass +class _Endpoint: + api_version: str = "2025-01-01" + service: str = "identity" + host: str = "identity.volcengineapi.com" + region: str = "cn-beijing" + + +def _patch_identity_request(monkeypatch, response): + captured = {} + + def _fake_resolve_credentials(service): + captured["credential_service"] = service + return _Creds() + + def _fake_resolve_endpoint(service): + captured["endpoint_service"] = service + return _Endpoint() + + def _fake_ve_request(**kwargs): + captured["request"] = kwargs + return response + + monkeypatch.setattr(identity_auth, "resolve_credentials", _fake_resolve_credentials) + monkeypatch.setattr(identity_auth, "resolve_endpoint", _fake_resolve_endpoint) + monkeypatch.setattr(identity_auth, "ve_request", _fake_ve_request) + return captured + + +def test_requires_api_key_passes_session_token(monkeypatch): + captured = _patch_identity_request( + monkeypatch, {"ResponseMetadata": {}, "Result": {"ApiKey": "api-key"}} + ) + + @requires_api_key(provider_name="provider") + def _call(api_key=None): + return api_key + + assert _call() == "api-key" + request = captured["request"] + assert captured["credential_service"] == "identity" + assert captured["endpoint_service"] == "identity" + assert request["session_token"] == "STS_TOKEN" + assert request["ak"] == "AK" + assert request["sk"] == "SK" + + +def test_requires_api_key_passes_session_token_for_async_function(monkeypatch): + captured = _patch_identity_request( + monkeypatch, {"ResponseMetadata": {}, "Result": {"ApiKey": "api-key"}} + ) + + @requires_api_key(provider_name="provider") + async def _call(api_key=None): + return api_key + + assert asyncio.run(_call()) == "api-key" + assert captured["request"]["session_token"] == "STS_TOKEN" + + +def test_requires_api_key_raises_domain_error_when_api_key_missing(monkeypatch): + _patch_identity_request(monkeypatch, {"ResponseMetadata": {}, "Result": {}}) + + @requires_api_key(provider_name="provider") + def _call(api_key=None): + return api_key + + with pytest.raises(ApiError, match="GetResourceApiKey did not return an ApiKey"): + _call() diff --git a/tests/toolkit/test_byteplus_single_region.py b/tests/toolkit/test_byteplus_single_region.py index 270a966..9372550 100644 --- a/tests/toolkit/test_byteplus_single_region.py +++ b/tests/toolkit/test_byteplus_single_region.py @@ -16,6 +16,7 @@ from agentkit.toolkit.config.choice_resolvers import resolve_field_choices from agentkit.toolkit.config.strategy_configs import CloudStrategyConfig +from agentkit.toolkit.volcengine.code_pipeline import VeCodePipeline from agentkit.toolkit.volcengine.cr import VeCR from agentkit.toolkit.volcengine.services.cr_service import CRServiceConfig @@ -37,6 +38,18 @@ def test_vecr_default_region_is_none() -> None: assert sig.parameters["region"].default is None +def test_vecr_session_token_is_keyword_compatible() -> None: + sig = inspect.signature(VeCR.__init__) + parameters = list(sig.parameters) + assert parameters.index("session_token") > parameters.index("provider") + + +def test_code_pipeline_session_token_is_keyword_compatible() -> None: + sig = inspect.signature(VeCodePipeline.__init__) + parameters = list(sig.parameters) + assert parameters.index("session_token") > parameters.index("provider") + + def test_cr_service_config_default_region_is_empty() -> None: cfg = CRServiceConfig() assert cfg.region == "" diff --git a/tests/toolkit/volcengine/test_provider_injection.py b/tests/toolkit/volcengine/test_provider_injection.py index d9ba3b4..7c0001d 100644 --- a/tests/toolkit/volcengine/test_provider_injection.py +++ b/tests/toolkit/volcengine/test_provider_injection.py @@ -30,6 +30,28 @@ def test_code_pipeline_uses_explicit_provider_over_env(monkeypatch) -> None: assert cp.host.endswith(".byteplusapi.com") +def test_code_pipeline_passes_resolved_session_token(monkeypatch) -> None: + from agentkit.platform.provider import ENV_CLOUD_PROVIDER + import agentkit.toolkit.volcengine.code_pipeline as code_pipeline_mod + from agentkit.toolkit.volcengine.code_pipeline import VeCodePipeline + + captured = {} + + def _fake_ve_request(**kwargs): + captured.update(kwargs) + return {"ResponseMetadata": {}, "Result": {"Id": "workspace-id"}} + + monkeypatch.setenv(ENV_CLOUD_PROVIDER, "volcengine") + monkeypatch.setenv("VOLCENGINE_ACCESS_KEY", "AK") + monkeypatch.setenv("VOLCENGINE_SECRET_KEY", "SK") + monkeypatch.setenv("VOLCENGINE_SESSION_TOKEN", "STS_TOKEN") + monkeypatch.setattr(code_pipeline_mod, "ve_request", _fake_ve_request) + + cp = VeCodePipeline(region="cn-beijing") + assert cp._get_default_workspace() == "workspace-id" + assert captured["session_token"] == "STS_TOKEN" + + def test_cr_uses_explicit_provider_over_env(monkeypatch) -> None: from agentkit.platform.provider import ENV_CLOUD_PROVIDER from agentkit.toolkit.volcengine.cr import VeCR @@ -45,6 +67,27 @@ def test_cr_uses_explicit_provider_over_env(monkeypatch) -> None: assert cr.host.endswith(".byteplusapi.com") +def test_cr_service_passes_resolved_session_token(monkeypatch) -> None: + from agentkit.platform.provider import ENV_CLOUD_PROVIDER + import agentkit.toolkit.volcengine.services.cr_service as cr_service_mod + from agentkit.toolkit.volcengine.services.cr_service import CRService + + captured = {} + + class _FakeVeCR: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setenv(ENV_CLOUD_PROVIDER, "volcengine") + monkeypatch.setenv("VOLCENGINE_ACCESS_KEY", "AK") + monkeypatch.setenv("VOLCENGINE_SECRET_KEY", "SK") + monkeypatch.setenv("VOLCENGINE_SESSION_TOKEN", "STS_TOKEN") + monkeypatch.setattr(cr_service_mod.ve_cr, "VeCR", _FakeVeCR) + + CRService() + assert captured["session_token"] == "STS_TOKEN" + + def test_tos_service_uses_explicit_provider_over_env(monkeypatch) -> None: from agentkit.platform.provider import ENV_CLOUD_PROVIDER import types diff --git a/tests/utils/test_engineering_standards.py b/tests/utils/test_engineering_standards.py index a97bce4..95c4593 100644 --- a/tests/utils/test_engineering_standards.py +++ b/tests/utils/test_engineering_standards.py @@ -109,6 +109,78 @@ def test_signed_request_retries_zero_means_single_attempt(monkeypatch, no_sleep) assert counter["attempts"] == 1 +def test_ve_request_signs_session_token(monkeypatch): + captured = {} + + class _FakeResponse: + status_code = 200 + headers = {} + + def json(self): + return {"ResponseMetadata": {}} + + def _fake_signed_request(method, url, headers, params, data): + captured.update( + { + "method": method, + "url": url, + "headers": headers, + "params": params, + "data": data, + } + ) + return _FakeResponse() + + monkeypatch.setattr(ve_sign, "_signed_request", _fake_signed_request) + + ve_sign.ve_request( + request_body={}, + action="ListRegistries", + ak="AK", + sk="SK", + service="cr", + version="2022-05-12", + region="cn-beijing", + host="cr.volcengineapi.com", + session_token="STS_TOKEN", + ) + + assert captured["headers"]["X-Security-Token"] == "STS_TOKEN" + assert "x-security-token" in captured["headers"]["Authorization"] + + +def test_ve_request_signs_session_token_from_existing_header(monkeypatch): + captured = {} + + class _FakeResponse: + status_code = 200 + headers = {} + + def json(self): + return {"ResponseMetadata": {}} + + def _fake_signed_request(method, url, headers, params, data): + captured["headers"] = headers + return _FakeResponse() + + monkeypatch.setattr(ve_sign, "_signed_request", _fake_signed_request) + + ve_sign.ve_request( + request_body={}, + action="GetResourceApiKey", + ak="AK", + sk="SK", + service="identity", + version="2025-01-01", + region="cn-beijing", + host="identity.volcengineapi.com", + header={"X-Security-Token": "HEADER_TOKEN"}, + ) + + assert captured["headers"]["X-Security-Token"] == "HEADER_TOKEN" + assert "x-security-token" in captured["headers"]["Authorization"] + + # --------------------------------------------------------------------------- # # http_defaults env clamping # --------------------------------------------------------------------------- #