Setting the file. One moment. 81
Creating Amazon Aurora Db Cluster With Instances
104
Routing Traffic With Route53 And CloudFront
Resilience Program Design
Creating API Gateway Stage
67
def test_mixed_models
— line 67
This file
- Number
- 30.27
- Position
- 27 of 35
- Type
- Python
- Size
- 7 KB
- Lines
- 170
scripts/test_iam_policy.py
Python·170 lines·7 KB
foundation_model_arn,
12 generate_policy,
13 is_mantle_model,
14 mantle_project_arn,
15 inference_profile_arn,
16 is_inference_profile,
17)
18
19REGION = "us-east-1"
20ACCOUNT = "123456789012"
21
22
23class TestIsInferenceProfile:
24 def test_geo_prefix_us(self):
25 assert is_inference_profile("us.anthropic.claude-sonnet-5")
26
27 def test_geo_prefix_eu(self):
28 assert is_inference_profile("eu.anthropic.claude-haiku-4-5-20251001-v1:0")
29
30 def test_no_prefix(self):
31 assert not is_inference_profile("anthropic.claude-haiku-4-5-20251001-v1:0")
32
33 def test_amazon_model(self):
34 assert not is_inference_profile("amazon.nova-lite-v1:0")
35
36
37class TestArnGeneration:
38 def test_foundation_model_arn(self):
39 arn = foundation_model_arn("anthropic.claude-haiku-4-5-20251001-v1:0")
40 assert arn == "arn:aws:bedrock:*::foundation-model/anthropic.claude-haiku-4-5-20251001-v1:0"
41
42 def test_inference_profile_arn(self):
43 arn = inference_profile_arn("us.anthropic.claude-sonnet-5", REGION, ACCOUNT)
44 assert arn == f"arn:aws:bedrock:{REGION}:{ACCOUNT}:inference-profile/us.anthropic.claude-sonnet-5"
45
46
47class TestGeneratePolicy:
48 def test_single_foundation_model(self):
49 policy = generate_policy(["anthropic.claude-haiku-4-5-20251001-v1:0"], REGION, ACCOUNT)
50 assert policy["Version"] == "2012-10-17"
51 stmt = policy["Statement"][0]
52 assert stmt["Effect"] == "Allow"
53 assert "bedrock:InvokeModel" in stmt["Action"]
54 assert "bedrock:InvokeModelWithResponseStream" in stmt["Action"]
55 assert len(stmt["Resource"]) == 1
56 assert "foundation-model/anthropic.claude-haiku-4-5-20251001-v1:0" in stmt["Resource"][0]
57
58 def test_inference_profile_generates_dual_arns(self):
59 policy = generate_policy(["us.anthropic.claude-sonnet-5"], REGION, ACCOUNT)
60 resources = policy["Statement"][0]["Resource"]
61 assert len(resources) == 2
62 has_foundation = any("foundation-model/" in r for r in resources)
63 has_profile = any("inference-profile/" in r for r in resources)
64 assert has_foundation
65 assert has_profile
66
67 def test_mixed_models(self):
68 models = [
69 "us.anthropic.claude-sonnet-5",
70 "amazon.nova-lite-v1:0",
71 ]
72 policy = generate_policy(models, REGION, ACCOUNT)
73 resources = policy["Statement"][0]["Resource"]
74 assert any("nova-lite" in r for r in resources)
75 assert any("inference-profile/" in r for r in resources)
76
77 def test_deduplicates_models(self):
78 models = ["amazon.nova-lite-v1:0", "amazon.nova-lite-v1:0"]
79 policy = generate_policy(models, REGION, ACCOUNT)
80 resources = policy["Statement"][0]["Resource"]
81 assert len(resources) == 1
82
83 def test_empty_models(self):
84 policy = generate_policy([], REGION, ACCOUNT)
85 assert policy["Statement"][0]["Resource"] == []
86
87
88class TestCLI:
89 def test_stdout_output(self):
90 result = subprocess.run( # nosec B603
91 [
92 sys.executable, str(Path(__file__).parent / "iam_policy.py"),
93 "--models", "amazon.nova-lite-v1:0",
94 "--region", REGION,
95 "--account-id", ACCOUNT,
96 ],
97 capture_output=True, text=True,
98 )
99 assert result.returncode == 0
100 policy = json.loads(result.stdout)
101 assert policy["Statement"][0]["Sid"] == "BedrockInvokeModelScoped"
102
103 def test_file_output(self, tmp_path):
104 out = tmp_path / "policy.json"
105 result = subprocess.run( # nosec B603
106 [
107 sys.executable, str(Path(__file__).parent / "iam_policy.py"),
108 "--models", "us.anthropic.claude-sonnet-5,amazon.nova-lite-v1:0",
109 "--region", REGION,
110 "--account-id", ACCOUNT,
111 "--output", str(out),
112 ],
113 capture_output=True, text=True,
114 )
115 assert result.returncode == 0
116 policy = json.loads(out.read_text())
117 assert len(policy["Statement"][0]["Resource"]) == 3
118
119
120def test_mantle_model_detection_excludes_gpt_oss():
121 assert is_mantle_model("openai.gpt-5.6-terra") is True
122 assert is_mantle_model("openai.gpt-5.5") is True
123 assert is_mantle_model("openai.gpt-oss-120b-1:0") is False
124 assert is_mantle_model("anthropic.claude-sonnet-4-6") is False
125
126
127def test_mantle_target_gets_mantle_actions_not_invoke_model():
128 # Regression: the generated policy emitted only bedrock:InvokeModel against a
129 # foundation-model ARN, which cannot authorize a mantle call — the migrated app
130 # would deploy and immediately fail authorization.
131 pol = generate_policy(["openai.gpt-5.6-terra"], "us-east-1", "111122223333")
132 sids = {s["Sid"] for s in pol["Statement"]}
133 assert "BedrockMantleInference" in sids
134 assert "BedrockMantleCallWithBearerToken" in sids
135 # An all-mantle run must NOT emit an InvokeModel statement with an empty
136 # Resource list — that is an invalid IAM policy.
137 assert "BedrockInvokeModelScoped" not in sids
138 for s in pol["Statement"]:
139 assert s.get("Resource"), s["Sid"]
140
141
142def test_mantle_bearer_token_is_scoped_to_star():
143 # AWS does not support narrowing CallWithBearerToken; it must be "*".
144 pol = generate_policy(["openai.gpt-5.5"], "us-east-1", "111122223333")
145 bearer = [s for s in pol["Statement"] if s["Sid"] == "BedrockMantleCallWithBearerToken"][0]
146 assert bearer["Resource"] == "*"
147 assert bearer["Action"] == ["bedrock-mantle:CallWithBearerToken"]
148
149
150def test_mantle_inference_scoped_to_account_and_region():
151 pol = generate_policy(["openai.gpt-5.5"], "us-west-2", "999988887777")
152 inf = [s for s in pol["Statement"] if s["Sid"] == "BedrockMantleInference"][0]
153 assert inf["Resource"] == "arn:aws:bedrock-mantle:us-west-2:999988887777:project/*"
154
155
156def test_mixed_targets_emit_both_runtime_and_mantle_statements():
157 pol = generate_policy(
158 ["openai.gpt-5.6-luna", "us.anthropic.claude-sonnet-4-6"], "us-east-1", "111122223333")
159 sids = [s["Sid"] for s in pol["Statement"]]
160 assert "BedrockInvokeModelScoped" in sids
161 assert "BedrockMantleInference" in sids
162 invoke = [s for s in pol["Statement"] if s["Sid"] == "BedrockInvokeModelScoped"][0]
163 # The mantle model must not leak into the InvokeModel resource list.
164 assert not any("gpt-5" in r for r in invoke["Resource"])
165
166
167def test_gpt_oss_stays_on_the_runtime_path():
168 pol = generate_policy(["openai.gpt-oss-120b-1:0"], "us-east-1", "111122223333")
169 sids = {s["Sid"] for s in pol["Statement"]}
170 assert sids == {"BedrockInvokeModelScoped"}