Skip to content

Commit 39108c6

Browse files
msyycmsyyc
committed
test case
1 parent 36a2676 commit 39108c6

File tree

6 files changed

+280
-39
lines changed

6 files changed

+280
-39
lines changed

‎sdk/resources/azure-mgmt-resource/azure/mgmt/resource/resources/_resource_management_client.py

Lines changed: 29 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,15 @@
66
# Changes may cause incorrect behavior and will be lost if the code is regenerated.
77
# --------------------------------------------------------------------------
88

9-
from copy import deepcopy
9+
from copy import deepcopy, copy
1010
from typing import Any, Optional, TYPE_CHECKING, cast
1111
from typing_extensions import Self
1212

1313
from azure.core.pipeline import policies
1414
from azure.core.rest import HttpRequest, HttpResponse
1515
from azure.core.settings import settings
1616
from azure.mgmt.core import ARMPipelineClient
17-
from azure.mgmt.core.policies import ARMAutoResourceProviderRegistrationPolicy
17+
from azure.mgmt.core.policies import ARMAutoResourceProviderRegistrationPolicy, PolicyTokenHeaderPolicy
1818
from azure.mgmt.core.tools import get_arm_endpoints
1919

2020
from . import models as _models
@@ -74,24 +74,35 @@ def __init__(
7474
credential=credential, subscription_id=subscription_id, credential_scopes=credential_scopes, **kwargs
7575
)
7676

77+
_policies_part1 = [
78+
policies.RequestIdPolicy(**kwargs),
79+
self._config.headers_policy,
80+
self._config.user_agent_policy,
81+
self._config.proxy_policy,
82+
policies.ContentDecodePolicy(**kwargs),
83+
ARMAutoResourceProviderRegistrationPolicy(),
84+
self._config.redirect_policy,
85+
self._config.retry_policy,
86+
self._config.authentication_policy,
87+
]
88+
_policies_part2 = [
89+
self._config.logging_policy,
90+
policies.DistributedTracingPolicy(**kwargs),
91+
policies.SensitiveHeaderCleanupPolicy(**kwargs) if self._config.redirect_policy else None,
92+
self._config.http_logging_policy,
93+
]
94+
_internal_policies = [copy(p) for p in _policies_part1 + _policies_part2]
95+
96+
_internal_client: ARMPipelineClient = ARMPipelineClient(
97+
base_url=cast(str, base_url), policies=_internal_policies, **kwargs
98+
)
7799
_policies = kwargs.pop("policies", None)
78100
if _policies is None:
79-
_policies = [
80-
policies.RequestIdPolicy(**kwargs),
81-
self._config.headers_policy,
82-
self._config.user_agent_policy,
83-
self._config.proxy_policy,
84-
policies.ContentDecodePolicy(**kwargs),
85-
ARMAutoResourceProviderRegistrationPolicy(),
86-
self._config.redirect_policy,
87-
self._config.retry_policy,
88-
self._config.authentication_policy,
89-
self._config.custom_hook_policy,
90-
self._config.logging_policy,
91-
policies.DistributedTracingPolicy(**kwargs),
92-
policies.SensitiveHeaderCleanupPolicy(**kwargs) if self._config.redirect_policy else None,
93-
self._config.http_logging_policy,
94-
]
101+
_policies = (
102+
_policies_part1
103+
+ [PolicyTokenHeaderPolicy(_internal_client, **kwargs), self._config.custom_hook_policy]
104+
+ _policies_part2
105+
)
95106
self._client: ARMPipelineClient = ARMPipelineClient(base_url=cast(str, base_url), policies=_policies, **kwargs)
96107

97108
client_models = {k: v for k, v in _models.__dict__.items() if isinstance(v, type)}

‎sdk/resources/azure-mgmt-resource/azure/mgmt/resource/resources/aio/_resource_management_client.py

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,15 @@
66
# Changes may cause incorrect behavior and will be lost if the code is regenerated.
77
# --------------------------------------------------------------------------
88

9-
from copy import deepcopy
9+
from copy import deepcopy, copy
1010
from typing import Any, Awaitable, Optional, TYPE_CHECKING, cast
1111
from typing_extensions import Self
1212

1313
from azure.core.pipeline import policies
1414
from azure.core.rest import AsyncHttpResponse, HttpRequest
1515
from azure.core.settings import settings
1616
from azure.mgmt.core import AsyncARMPipelineClient
17-
from azure.mgmt.core.policies import AsyncARMAutoResourceProviderRegistrationPolicy
17+
from azure.mgmt.core.policies import AsyncARMAutoResourceProviderRegistrationPolicy, AsyncPolicyTokenHeaderPolicy
1818
from azure.mgmt.core.tools import get_arm_endpoints
1919

2020
from .. import models as _models
@@ -73,25 +73,38 @@ def __init__(
7373
self._config = ResourceManagementClientConfiguration(
7474
credential=credential, subscription_id=subscription_id, credential_scopes=credential_scopes, **kwargs
7575
)
76+
_policies_part1 = [
77+
policies.RequestIdPolicy(**kwargs),
78+
self._config.headers_policy,
79+
self._config.user_agent_policy,
80+
self._config.proxy_policy,
81+
policies.ContentDecodePolicy(**kwargs),
82+
AsyncARMAutoResourceProviderRegistrationPolicy(),
83+
self._config.redirect_policy,
84+
self._config.retry_policy,
85+
self._config.authentication_policy,
86+
]
87+
_policies_part2 = [
88+
self._config.logging_policy,
89+
policies.DistributedTracingPolicy(**kwargs),
90+
policies.SensitiveHeaderCleanupPolicy(**kwargs) if self._config.redirect_policy else None,
91+
self._config.http_logging_policy,
92+
]
93+
_internal_polices = [copy(p) for p in _policies_part1 + _policies_part2]
94+
_internal_client: AsyncARMPipelineClient = AsyncARMPipelineClient(
95+
base_url=cast(str, base_url), policies=_internal_polices, **kwargs
96+
)
7697

7798
_policies = kwargs.pop("policies", None)
7899
if _policies is None:
79-
_policies = [
80-
policies.RequestIdPolicy(**kwargs),
81-
self._config.headers_policy,
82-
self._config.user_agent_policy,
83-
self._config.proxy_policy,
84-
policies.ContentDecodePolicy(**kwargs),
85-
AsyncARMAutoResourceProviderRegistrationPolicy(),
86-
self._config.redirect_policy,
87-
self._config.retry_policy,
88-
self._config.authentication_policy,
89-
self._config.custom_hook_policy,
90-
self._config.logging_policy,
91-
policies.DistributedTracingPolicy(**kwargs),
92-
policies.SensitiveHeaderCleanupPolicy(**kwargs) if self._config.redirect_policy else None,
93-
self._config.http_logging_policy,
94-
]
100+
_policies = (
101+
_policies_part1
102+
+ [
103+
AsyncPolicyTokenHeaderPolicy(_internal_client, **kwargs),
104+
self._config.custom_hook_policy,
105+
]
106+
+ _policies_part2
107+
)
95108
self._client: AsyncARMPipelineClient = AsyncARMPipelineClient(
96109
base_url=cast(str, base_url), policies=_policies, **kwargs
97110
)

‎sdk/resources/azure-mgmt-resource/samples/policy_token_header_policy.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# pylint: disable=line-too-long,useless-suppression
2+
# coding=utf-8
3+
# --------------------------------------------------------------------------
4+
# Copyright (c) Microsoft Corporation. All rights reserved.
5+
# Licensed under the MIT License. See License.txt in the project root for license information.
6+
# Code generated by Microsoft (R) AutoRest Code Generator.
7+
# Changes may cause incorrect behavior and will be lost if the code is regenerated.
8+
# --------------------------------------------------------------------------
9+
10+
import os
11+
from azure.identity import DefaultAzureCredential
12+
from dotenv import load_dotenv
13+
from azure.mgmt.resource.resources import ResourceManagementClient
14+
15+
16+
import sys
17+
import logging
18+
19+
logging.basicConfig(
20+
level=logging.DEBUG, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", stream=sys.stdout
21+
)
22+
23+
24+
def check_policy_token_header_exist(request):
25+
if request.http_request.method != "GET":
26+
assert "x-ms-policy-external-evaluations" in request.http_request.headers
27+
28+
29+
load_dotenv()
30+
client = ResourceManagementClient(
31+
credential=DefaultAzureCredential(), subscription_id=os.getenv("AZURE_SUBSCRIPTION_ID", ""), logging_enable=True
32+
)
33+
34+
# resource_group_name="test-resource-group-yyc5"
35+
# response = client.resource_groups.create_or_update(
36+
# resource_group_name=resource_group_name,
37+
# parameters={"location": "eastus"},
38+
# # raw_request_hook=check_policy_token_header_exist,
39+
# # acquire_policy_token=True,
40+
# )
41+
# print(response)
42+
43+
client.resources.check_existence_by_id(
44+
resource_id="test",
45+
api_version="2021-04-01",
46+
acquire_policy_token=True,
47+
)

‎sdk/resources/azure-mgmt-resource/setup.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,11 @@
2222

2323
# Version extraction inspired from 'requests'
2424
with open(
25-
os.path.join(package_folder_path, "version.py")
26-
if os.path.exists(os.path.join(package_folder_path, "version.py"))
27-
else os.path.join(package_folder_path, "_version.py"),
25+
(
26+
os.path.join(package_folder_path, "version.py")
27+
if os.path.exists(os.path.join(package_folder_path, "version.py"))
28+
else os.path.join(package_folder_path, "_version.py")
29+
),
2830
"r",
2931
) as fd:
3032
version = re.search(r'^VERSION\s*=\s*[\'"]([^\'"]*)[\'"]', fd.read(), re.MULTILINE).group(1)

‎sdk/resources/azure-mgmt-resource/tests/test_policy_token_header.py

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
# coding=utf-8
2+
# --------------------------------------------------------------------------
3+
# Copyright (c) Microsoft Corporation. All rights reserved.
4+
# Licensed under the MIT License. See License.txt in the project root for license information.
5+
# Code generated by Microsoft (R) AutoRest Code Generator.
6+
# Changes may cause incorrect behavior and will be lost if the code is regenerated.
7+
# --------------------------------------------------------------------------
8+
import pytest
9+
from azure.mgmt.resource.resources import ResourceManagementClient
10+
from azure.core.exceptions import HttpResponseError
11+
from devtools_testutils import AzureMgmtRecordedTestCase, recorded_by_proxy, RandomNameResourceGroupPreparer
12+
13+
AZURE_LOCATION = "eastus"
14+
15+
16+
def check_policy_token_header_exist_lro(request):
17+
if request.http_request.method != "GET":
18+
check_policy_token_header_exist(request)
19+
20+
21+
def check_policy_token_header_exist(request):
22+
assert "x-ms-policy-external-evaluations" in request.http_request.headers
23+
24+
25+
def check_policy_token_header_not_exist(request):
26+
assert "x-ms-policy-external-evaluations" not in request.http_request.headers
27+
28+
29+
@pytest.mark.live_test_only
30+
class TestPolicyTokenHeader(AzureMgmtRecordedTestCase):
31+
def setup_method(self, method):
32+
self.client = self.create_mgmt_client(ResourceManagementClient)
33+
34+
@RandomNameResourceGroupPreparer(location=AZURE_LOCATION)
35+
@recorded_by_proxy
36+
def test_resource_groups_check_existence(self, resource_group):
37+
with pytest.raises(HttpResponseError, match="status code is"):
38+
self.client.resource_groups.check_existence(
39+
resource_group_name=resource_group.name,
40+
acquire_policy_token=True,
41+
raw_request_hook=check_policy_token_header_exist,
42+
)
43+
44+
@RandomNameResourceGroupPreparer(location=AZURE_LOCATION)
45+
@recorded_by_proxy
46+
def test_resource_groups_create_or_update(self, resource_group):
47+
self.client.resource_groups.create_or_update(
48+
resource_group_name=resource_group.name,
49+
parameters={
50+
"location": AZURE_LOCATION,
51+
},
52+
raw_request_hook=check_policy_token_header_not_exist,
53+
)
54+
55+
@RandomNameResourceGroupPreparer(location=AZURE_LOCATION)
56+
@recorded_by_proxy
57+
def test_resource_groups_delete(self, resource_group):
58+
self.client.resource_groups.begin_delete(
59+
resource_group_name=resource_group.name,
60+
acquire_policy_token=True,
61+
raw_request_hook=check_policy_token_header_exist_lro,
62+
).result()
63+
64+
@RandomNameResourceGroupPreparer(location=AZURE_LOCATION)
65+
@recorded_by_proxy
66+
def test_resource_groups_get(self, resource_group):
67+
# policy token header shall not exist for GET operation
68+
self.client.resource_groups.get(
69+
resource_group_name=resource_group.name,
70+
acquire_policy_token=True,
71+
raw_request_hook=check_policy_token_header_not_exist,
72+
)
73+
74+
@recorded_by_proxy
75+
def test_resource_check_existence_by_id(self):
76+
# If path path contains no subscriptionId, it shall raise error
77+
with pytest.raises(HttpResponseError, match="Failed to get subscriptionId from request url"):
78+
self.client.resources.check_existence_by_id(
79+
resource_id="test",
80+
api_version="2021-04-01",
81+
acquire_policy_token=True,
82+
)

‎sdk/resources/azure-mgmt-resource/tests/test_policy_token_header_async.py

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
# coding=utf-8
2+
# --------------------------------------------------------------------------
3+
# Copyright (c) Microsoft Corporation. All rights reserved.
4+
# Licensed under the MIT License. See License.txt in the project root for license information.
5+
# Code generated by Microsoft (R) AutoRest Code Generator.
6+
# Changes may cause incorrect behavior and will be lost if the code is regenerated.
7+
# --------------------------------------------------------------------------
8+
import pytest
9+
from azure.mgmt.resource.resources.aio import ResourceManagementClient
10+
from azure.core.exceptions import HttpResponseError
11+
from devtools_testutils import AzureMgmtRecordedTestCase, RandomNameResourceGroupPreparer
12+
from devtools_testutils.aio import recorded_by_proxy_async
13+
14+
AZURE_LOCATION = "eastus"
15+
16+
17+
def check_policy_token_header_exist_lro(request):
18+
if request.http_request.method != "GET":
19+
check_policy_token_header_exist(request)
20+
21+
22+
def check_policy_token_header_exist(request):
23+
assert "x-ms-policy-external-evaluations" in request.http_request.headers
24+
25+
26+
def check_policy_token_header_not_exist(request):
27+
assert "x-ms-policy-external-evaluations" not in request.http_request.headers
28+
29+
30+
@pytest.mark.live_test_only
31+
class TestPolicyTokenHeaderAsync(AzureMgmtRecordedTestCase):
32+
def setup_method(self, method):
33+
self.client = self.create_mgmt_client(ResourceManagementClient, is_async=True)
34+
35+
@RandomNameResourceGroupPreparer(location=AZURE_LOCATION)
36+
@recorded_by_proxy_async
37+
async def test_resource_groups_check_existence(self, resource_group):
38+
with pytest.raises(HttpResponseError, match="status code is"):
39+
await self.client.resource_groups.check_existence(
40+
resource_group_name=resource_group.name,
41+
acquire_policy_token=True,
42+
raw_request_hook=check_policy_token_header_exist,
43+
)
44+
45+
@RandomNameResourceGroupPreparer(location=AZURE_LOCATION)
46+
@recorded_by_proxy_async
47+
async def test_resource_groups_create_or_update(self, resource_group):
48+
await self.client.resource_groups.create_or_update(
49+
resource_group_name=resource_group.name,
50+
parameters={
51+
"location": AZURE_LOCATION,
52+
},
53+
raw_request_hook=check_policy_token_header_not_exist,
54+
)
55+
56+
@RandomNameResourceGroupPreparer(location=AZURE_LOCATION)
57+
@recorded_by_proxy_async
58+
async def test_resource_groups_delete(self, resource_group):
59+
await (
60+
await self.client.resource_groups.begin_delete(
61+
resource_group_name=resource_group.name,
62+
acquire_policy_token=True,
63+
raw_request_hook=check_policy_token_header_exist_lro,
64+
)
65+
).result()
66+
67+
@RandomNameResourceGroupPreparer(location=AZURE_LOCATION)
68+
@recorded_by_proxy_async
69+
async def test_resource_groups_get(self, resource_group):
70+
# policy token header shall not exist for GET operation
71+
await self.client.resource_groups.get(
72+
resource_group_name=resource_group.name,
73+
acquire_policy_token=True,
74+
raw_request_hook=check_policy_token_header_not_exist,
75+
)
76+
77+
@pytest.mark.asyncio
78+
@recorded_by_proxy_async
79+
async def test_resource_check_existence_by_id(self):
80+
# If path path contains no subscriptionId, it shall raise error
81+
with pytest.raises(HttpResponseError, match="Failed to get subscriptionId from request url"):
82+
await self.client.resources.check_existence_by_id(
83+
resource_id="test",
84+
api_version="2021-04-01",
85+
acquire_policy_token=True,
86+
)

0 commit comments

Comments
 (0)