diff --git a/features/steps/acms_plan_execution_integration_steps.py b/features/steps/acms_plan_execution_integration_steps.py
index e7042d8f1..8cf5f3dc3 100644
--- a/features/steps/acms_plan_execution_integration_steps.py
+++ b/features/steps/acms_plan_execution_integration_steps.py
@@ -29,6 +29,14 @@ from cleveragents.application.services.plan_executor import (
 from cleveragents.tool.runner import ToolRunner
 
 
+@when("I prepare LLM context$")
+def step_prepare_llm_context_generic(context: Any) -> None:
+    """Prepare LLM context. Handles both plain 'prepare LLM context' and qualified versions."""
+    if context.integration is not None:
+        raw = getattr(context, "raw_context", {"file_type": "python"})
+        context.llm_context = context.integration.prepare_llm_context(raw)
+
+
 @given("I have a plan execution ACMS integration")
 def step_have_integration(context: Any) -> None:
     """Initialize a plan execution ACMS integration."""
@@ -232,6 +240,30 @@ def step_have_scope_rules(context: Any) -> None:
     )
 
 
+@given("policy{count:d} has priority_weight {weight:f}")
+def step_set_policy_priority(context: Any, count: int, weight: float) -> None:
+    """Set the priority weight for a specific policy by index."""
+    idx = count - 1
+    if context.policy_config and context.policy_config.policies:
+        if 0 <= idx < len(context.policy_config.policies):
+            context.policy_config.policies[idx].priority_weight = float(weight)
+            # Rebuild integration with updated config
+            context.integration = PlanExecutionACMSIntegration(
+                policy_config=context.policy_config,
+            )
+
+
+@given("the policy has budget_override {amount:int}")
+def step_set_policy_budget(context: Any, amount: int) -> None:
+    """Set the budget override for the first policy in the configuration."""
+    if context.policy_config and context.policy_config.policies:
+        context.policy_config.policies[0].budget_override = int(amount)
+        # Rebuild integration with updated config
+        context.integration = PlanExecutionACMSIntegration(
+            policy_config=context.policy_config,
+        )
+
+
 @given("the scope rule is file_type equals {value}")
 def step_scope_rule(context: Any, value: str) -> None:
     """Set the scope rule."""
@@ -259,6 +291,31 @@ def step_policy_applied(context: Any) -> None:
     assert "policy1" in policies_applied
 
 
+@then("the assembled context should have budget {amount:int}")
+def step_check_budget(context: Any, amount: int) -> None:
+    """Check that the assembled context includes the expected budget override."""
+    assert context.llm_context is not None
+    assembled_data = context.llm_context.get("assembled_data", {})
+    assert "budget" in assembled_data, (
+        f"Expected 'budget' key in assembled_data, got keys: {list(assembled_data.keys())}"
+    )
+    assert assembled_data["budget"] == int(amount)
+
+
+@then("the policy should not be applied")
+def step_policy_not_applied(context: Any) -> None:
+    """Check that the policy was NOT applied."""
+    assert context.llm_context is not None
+    policies_applied = context.llm_context.get("policies_applied", [])
+    assert "policy1" not in policies_applied
+
+    # Also verify budget is absent when no matching policy applies
+    assembled_data = context.llm_context.get("assembled_data", {})
+    assert "budget" not in assembled_data, (
+        f"Expected no 'budget' key when scope doesn't match, got keys: {list(assembled_data.keys())}"
+    )
+
+
 @given("I have a plan execution ACMS integration with a policy configuration")
 def step_have_integration_with_policy(context: Any) -> None:
     """Create a plan execution ACMS integration with a policy configuration."""