Add automatic test generation after code generation

Introduces a new workflow that generates tests for newly created Azure CLI modules/extensions using LLM sampling. Updates README to reflect new features, adds the generate_tests function to helpers.py, and integrates test generation into the main code generation flow in main.py.

Author: a0x1ab

tools/aaz-flow/README.md

@@ -8,5 +8,7 @@ Please note that AAZ Flow is currently in early development. The functionality a
 
 ## Implemetation
 1. Performs elicitation for user input to perform workflow
-2. Generates elicitation questions using llm prompts
-3. Executes AAZ Flow commands directly
+2. Generates content using llm sampling
+3. Executes AAZ Flow commands directly
+4. Generates tests using llm sampling
+5. Uses tool transformation to make the internal tooling more friendly for llms

tools/aaz-flow/helpers.py

@@ -1,4 +1,5 @@
 import asyncio
+from pathlib import Path
 import sys
 from typing import Literal
 from fastmcp import Context
@@ -7,6 +8,13 @@
 import subprocess as sp
 from models import AAZRequest
 
+paths = {
+    "aaz": os.getenv("AAZ_PATH", "/workspaces-src/aaz"),
+    "cli": "/workspaces/azure-cli",
+    "cli_extension": os.getenv("CLI_EXTENSION_PATH", "/workspaces-src/azure-cli-extensions"),
+    "swagger_path": os.getenv("SWAGGER_PATH", "/workspaces-src/azure-rest-api-specs")
+}
+
 async def fetch_available_services():
     """Fetch available services from azure-rest-api-specs repository."""
     url = "https://api.github.com/repos/a0x1ab/azure-rest-api-specs/contents/specification"
@@ -22,12 +30,6 @@ async def fetch_available_services():
 
 async def validate_paths(ctx: Context) -> dict:
     """Validate and get correct paths for required directories."""
-    paths = {
-        "aaz": os.getenv("AAZ_PATH", "/workspaces-src/aaz"),
-        "cli": "/workspaces/azure-cli",
-        "cli_extension": os.getenv("CLI_EXTENSION_PATH", "/workspaces-src/azure-cli-extensions"),
-        "swagger_path": os.getenv("SWAGGER_PATH", "/workspaces-src/azure-rest-api-specs")
-    }
 
     await ctx.info("az_cli : Validating local paths...")
     await ctx.report_progress(progress=5, total=100)
@@ -182,7 +184,6 @@ async def browse_specs(ctx: Context, base_path: str):
             return result
 
 async def run_command(ctx: Context, command: str, step_name: str, progress_start: int, progress_end: int):
-    """Run a shell command and report progress."""
     await ctx.info(f"az_cli : Starting: {step_name}")
     process = await asyncio.create_subprocess_shell(
         command,
@@ -197,20 +198,23 @@ async def run_command(ctx: Context, command: str, step_name: str, progress_start
     while True:
         line = await process.stdout.readline()
         if not line:
-            break
+            if process.returncode is not None:
+                break
+            await asyncio.sleep(0.1)
+            continue
         lines_count += 1
         await ctx.info(f"az_cli : {line.decode().rstrip()}")
         progress = progress_start + min(progress_range, int((lines_count / total_lines_estimate) * progress_range))
         await ctx.report_progress(progress, 100)
 
     await process.wait()
+
     if process.returncode != 0:
         raise RuntimeError(f"{step_name} failed: {command}")
 
     await ctx.report_progress(progress_end, 100)
     await ctx.info(f"az_cli : Completed: {step_name}")
 
-
 async def execute_commands(ctx: Context, paths: dict, request: AAZRequest):
     cmd1 = (
         f"aaz-dev command-model generate-from-swagger "
@@ -241,3 +245,41 @@ async def execute_commands(ctx: Context, paths: dict, request: AAZRequest):
         return f"Code generation failed: {str(e)}"
 
     return "Azure CLI code generation completed successfully!"
+
+async def generate_tests(ctx: "Context"):
+    await ctx.info("Starting test generation workflow.")
+
+    module_name = getattr(ctx, "generated_module", None)
+    if not module_name:
+        response = await ctx.elicit("Enter the module/extension name to generate tests for:")
+        if not response.action == "accept" or not response.data:
+            return "Test generation cancelled."
+        module_name = response.data
+    else:
+        await ctx.info(f"Detected generated module: {module_name}")
+
+    aaz_path = Path(f"{paths['cli']}/src/azure-cli/azure/cli/command_modules/{module_name}/aaz")
+    if not aaz_path.exists():
+        return f"AAZ path not found for module '{module_name}'"
+
+    commands = []
+    for file in aaz_path.rglob("*.py"):
+        with open(file, "r", encoding="utf-8") as f:
+            for line in f:
+                if line.strip().startswith("def "):
+                    commands.append(line.strip().replace("def ", "").split("(")[0])
+
+    test_dir = Path(f"{paths['cli']}/src/azure-cli/azure/cli/command_modules/{module_name}/tests/latest")
+    test_dir.mkdir(parents=True, exist_ok=True)
+    test_file = test_dir / f"test_{module_name}.py"
+
+    with open(test_file, "w", encoding="utf-8") as f:
+        f.write("import unittest\n")
+        f.write("from azure.cli.testsdk import ScenarioTest\n\n")
+        f.write(f"class {module_name.capitalize()}ScenarioTest(ScenarioTest):\n\n")
+        for cmd in commands:
+            f.write(f"    def test_{cmd}(self):\n")
+            f.write(f"        self.cmd('az {module_name} {cmd} --resource-name test-resource')\n\n")
+
+    await ctx.info(f"Generated test file: {test_file}")
+    return f"Test generation completed for module '{module_name}'."

tools/aaz-flow/main.py

@@ -1,7 +1,8 @@
 import os
+from pathlib import Path
 from fastmcp import FastMCP, Context
 from models import AAZRequest
-from helpers import execute_commands, validate_paths, get_extension_name, get_swagger_config
+from helpers import generate_tests, execute_commands, validate_paths, get_extension_name, get_swagger_config
 
 mcp = FastMCP("AAZ Flow")
 
@@ -53,7 +54,18 @@ async def generate_code(ctx: Context):
 
     await execute_commands(ctx, paths, request)
     await ctx.report_progress(100, 100)
-    return f"Code generation completed for extension/module '{extension_name}'."
+    await ctx.info(f"Code generation completed for extension/module '{extension_name}'.")
+
+    ctx.generated_module = extension_name
+
+    await ctx.info("Automatically generating tests for the newly generated module...")
+    try:
+        test_result = await generate_tests(ctx)
+        await ctx.info(f"Automatic test generation result: {test_result}")
+    except Exception as e:
+        await ctx.info(f"Automatic test generation failed: {str(e)}")
+
+    return f"Code generation and test generation completed for extension/module '{extension_name}'."
 
 if __name__ == "__main__":
     mcp.run(transport="stdio")