added ideal template, improved robustness

Author: a0x1ab

tools/aaz-flow/prompt_templates.py

@@ -5,6 +5,32 @@
 
 """Prompt templates and static guidance for AAZ Flow tools."""
 
+IDEAL_STYLE = """
+    @ResourceGroupPreparer(name_prefix="cli_test_ag_with_non_v2_sku_", location="westus")
+    def test_ag_with_non_v2_sku(self):
+        self.kwargs.update({
+            "ag_name": self.create_random_name("ag-", 12),
+            "port_name": self.create_random_name("port-", 12),
+            "lisener_name": self.create_random_name("lisener-", 12),
+            "rule_name": self.create_random_name("rule-", 12),
+        })
+
+        self.cmd("network application-gateway create -n {ag_name} -g {rg} --sku WAF_Medium")
+
+        self.kwargs["front_ip"] = self.cmd("network application-gateway show -n {ag_name} -g {rg}").get_output_in_json()["frontendIPConfigurations"][0]["name"]
+        self.cmd("network application-gateway frontend-port create -n {port_name} -g {rg} --gateway-name {ag_name} --port 8080")
+        self.cmd("network application-gateway http-listener create -n {lisener_name} -g {rg} --gateway-name {ag_name} --frontend-ip {front_ip} --frontend-port {port_name}")
+
+        self.cmd(
+            "network application-gateway rule create -n {rule_name} -g {rg} --gateway-name {ag_name} --http-listener {lisener_name}",
+            checks=[
+                self.check("name", "{ag_name}"),
+                self.check("sku.tier", "WAF")
+            ]
+        )
+
+        self.cmd("network application-gateway delete -n {ag_name} -g {rg}")
+"""
 
 def get_testgen_static_instructions() -> str:
     return (
@@ -13,13 +39,10 @@ def get_testgen_static_instructions() -> str:
         "Generate tests that achieve at least 80%% coverage of methods and parameters covering primary commands for the target module.\n"
         "To understand the primary commands that need to be tested, read through and understand the target module's generated AAZ commands.\n"
         "Constraints: \n"
-        "- Include necessary imports: azure.cli.testsdk imports and others only as required and seen in reference.\n"
-        "- Use self.kwargs for dynamic values.\n"
-        "- Use ResourceGroupPreparer if a resource group is implied.\n"
-        "- Add minimal checks (e.g., self.check) where sensible.\n"
         "- Keep tests safe-by-default; avoid destructive operations unless clearly required.\n"
         "- Ensure tests can run in parallel without conflicts.\n"
         "- If tests are large and can be safely and logically split, create multiple test methods (i.e. avoid a single CRUD test if possible, split it into multiple tests if logically and safely separable).\n"
+        "- It is highly preferred that all CRUD operations are not coupled in a single test.\n"
         "- Output only valid Python code for the test file, nothing else."
     )
 

tools/aaz-flow/testgen.py

@@ -8,7 +8,7 @@
 import asyncio
 import random
 from fastmcp import Context
-from prompt_templates import get_testgen_static_instructions, REF_STYLE_LABEL
+from prompt_templates import get_testgen_static_instructions, REF_STYLE_LABEL, IDEAL_STYLE
 
 
 async def check_module_status(ctx: Context, paths: dict):
@@ -120,6 +120,8 @@ def build_testgen_prompt(
 
     if reference_snippet:
         parts.append(REF_STYLE_LABEL + reference_snippet)
+    
+    parts.append("Here is the ideal test style example to follow:\n" + IDEAL_STYLE)
 
     return "\n\n".join(parts)
 
@@ -147,6 +149,11 @@ def strip_code_fences(text: str) -> str:
         return max(blocks, key=len)
     return text.strip()
 
+def strip_shebang(text: str) -> str:
+    lines = text.strip().splitlines()
+    if lines and lines[0].startswith("#!"):
+        return "\n".join(lines[1:]).strip()
+    return text.strip()
 
 async def generate_tests(ctx: Context, paths: dict):
     module_name, commands, test_file = await check_path_status(ctx, paths)
@@ -171,7 +178,6 @@ async def generate_tests(ctx: Context, paths: dict):
     sampling_prompt = build_testgen_prompt(
         module_name, commands, reference_snippet, extracted_examples
     )
-    await ctx.info("Constructed test generation prompt as follows:\n" + sampling_prompt)
     max_retries = int(os.getenv("TESTGEN_RETRIES", "5"))
     base_delay = float(os.getenv("TESTGEN_RETRY_BASE_DELAY", "2"))
 
@@ -186,7 +192,7 @@ async def generate_tests(ctx: Context, paths: dict):
                 )
             sampled = await ctx.sample(sampling_prompt)
             raw_content = (getattr(sampled, "text", "") or "").strip()
-            content = strip_code_fences(raw_content)
+            content = strip_shebang(strip_code_fences(raw_content))
             if content:
                 break
             last_err = RuntimeError("Empty content returned from provider")