ARBuilder

Instructions

Implementation Reference

• src/mcp/tools/generate_tests.py:89-199 (handler)

  The main handler logic for the 'generate_tests' tool, implementing LLM-based test generation for Rust/Stylus contracts.

  def execute(
      self,
      contract_code: str,
      test_framework: str = "rust_native",
      test_types: Optional[list[str]] = None,
      coverage_focus: Optional[list[str]] = None,
      **kwargs,
  ) -> dict:
      """
      Generate tests for a Stylus contract.
  
      Args:
          contract_code: The contract code to generate tests for.
          test_framework: Test framework (rust_native, foundry).
          test_types: Types of tests (unit, integration, fuzz).
          coverage_focus: Specific functions to focus on.
  
      Returns:
          Dict with tests, test_summary, coverage_estimate,
          setup_instructions.
      """
      if not contract_code or not contract_code.strip():
          return {"error": "Contract code is required and cannot be empty"}
  
      contract_code = contract_code.strip()
      test_types = test_types or ["unit"]
  
      if not self._is_valid_contract(contract_code):
          return {
              "error": (
                  "Invalid contract code. Please provide valid"
                  " Stylus/Rust code with struct and impl blocks."
              ),
              "warnings": ["Could not parse contract structure"],
          }
  
      try:
          # Build LLM prompt
          system_prompt = (
              SYSTEM_PROMPT_FOUNDRY
              if test_framework == "foundry"
              else SYSTEM_PROMPT_RUST
          )
  
          focus_hint = ""
          if coverage_focus:
              focus_hint = (
                  "\n\nFocus test coverage on these functions: "
                  + ", ".join(coverage_focus)
              )
  
          type_hint = ""
          if "fuzz" in test_types:
              type_hint += (
                  "\n\nInclude fuzz/property-based tests using proptest."
              )
          if "integration" in test_types:
              type_hint += (
                  "\n\nInclude integration tests that test"
                  " multi-function workflows."
              )
  
          user_prompt = (
              f"Generate tests for this contract:"
              f"\n\n```rust\n{contract_code}\n```"
              f"{focus_hint}{type_hint}"
          )
  
          messages = [
              {"role": "system", "content": system_prompt},
              {"role": "user", "content": user_prompt},
          ]
  
          response = self._call_llm(
              messages=messages,
              temperature=0.2,
              max_tokens=8192,
          )
  
          # Extract test code from code block
          lang = "solidity" if test_framework == "foundry" else "rust"
          pattern = rf"```(?:{lang})\n([\s\S]*?)```"
          test_match = re.search(pattern, response)
          tests = test_match.group(1).strip() if test_match else response
  
          # Analyze contract for coverage stats
          contract_info = self._analyze_contract(contract_code)
  
          # Setup instructions
          setup = (
              self._get_foundry_setup()
              if test_framework == "foundry"
              else self._get_rust_setup()
          )
  
          # Analyze generated tests
          test_summary = self._generate_summary(tests, test_types)
          coverage_estimate = self._estimate_coverage(
              contract_info, tests
          )
  
          return {
              "tests": tests,
              "test_summary": test_summary,
              "coverage_estimate": coverage_estimate,
              "setup_instructions": setup,
          }
  
      except Exception as e:
          logger.exception("Test generation failed")
          return {"error": f"Test generation failed: {str(e)}"}

• src/mcp/server.py:207-236 (schema)

  Input schema definition for the 'generate_tests' MCP tool.

      "name": "generate_tests",
      "description": "Generate test cases for Stylus smart contracts.",
      "inputSchema": {
          "type": "object",
          "properties": {
              "contract_code": {
                  "type": "string",
                  "description": "The contract code to generate tests for",
              },
              "test_framework": {
                  "type": "string",
                  "enum": ["rust_native", "foundry", "hardhat"],
                  "description": "Test framework to use (default: rust_native)",
                  "default": "rust_native",
              },
              "test_types": {
                  "type": "array",
                  "items": {"type": "string", "enum": ["unit", "integration", "fuzz"]},
                  "description": 'Types of tests to generate (default: ["unit"])',
                  "default": ["unit"],
              },
              "coverage_focus": {
                  "type": "array",
                  "items": {"type": "string"},
                  "description": "Specific functions to focus on",
              },
          },
          "required": ["contract_code"],
      },
  },

• src/mcp/server.py:879-879 (registration)

  Registration of the 'generate_tests' tool in the MCPServer instance.

  "generate_tests": GenerateTestsTool(),