python-sdk/pkg/hanzo-mcp/test_exact_tools.py

#!/usr/bin/env python3
"""
Test Runner for Exact 6-Tool Implementation
===========================================

Test the exact tool specifications with real workspace scenarios.
"""

import json
import shutil
import asyncio
import tempfile
from pathlib import Path
from dataclasses import asdict

from hanzo_mcp.exact_tools import (
    FmtArgs,
    EditArgs,
    LintArgs,
    TestArgs,
    BuildArgs,
    GuardArgs,
    GuardRule,
    TargetSpec,
    tools,
)


async def test_go_workspace_scenario():
    """Test Go workspace with go.work file"""
    print("🔧 Testing Go workspace scenario...")
    
    # Create temporary Go workspace
    with tempfile.TemporaryDirectory() as temp_dir:
        workspace_dir = Path(temp_dir)
        
        # Create go.work file
        go_work = workspace_dir / "go.work"
        go_work.write_text("""go 1.21

use (
    ./api
    ./cli
)
""")
        
        # Create api module
        api_dir = workspace_dir / "api"
        api_dir.mkdir()
        
        api_go_mod = api_dir / "go.mod"
        api_go_mod.write_text("module github.com/luxfi/api\n\ngo 1.21\n")
        
        api_main = api_dir / "main.go"
        api_main.write_text("""package main

import (
	"fmt"
	"net/http"
)

func UserHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Hello User")
}

func main() {
	http.HandleFunc("/user", UserHandler)
	http.ListenAndServe(":8080", nil)
}
""")
        
        # Create cli module
        cli_dir = workspace_dir / "cli"
        cli_dir.mkdir()
        
        cli_go_mod = cli_dir / "go.mod"
        cli_go_mod.write_text("module github.com/luxfi/cli\n\ngo 1.21\n")
        
        cli_main = cli_dir / "main.go"
        cli_main.write_text("""package main

import "fmt"

func GreetUser(name string) {
	fmt.Printf("Hello %s\\n", name)
}

func main() {
	GreetUser("World")
}
""")
        
        # Test workspace detection
        target_spec = TargetSpec(
            target="ws",
            root=str(workspace_dir)
        )
        
        # Test fmt tool
        fmt_result = await tools.fmt(target_spec, FmtArgs(opts={"local_prefix": "github.com/luxfi"}))
        print(f"  fmt result: {'✅' if fmt_result.ok else '❌'} {fmt_result.language_used}")
        print(f"    root: {fmt_result.root}")
        print(f"    backend: {fmt_result.backend_used}")
        
        # Test edit tool - organize imports
        edit_result = await tools.edit(target_spec, EditArgs(op="organize_imports"))
        print(f"  edit result: {'✅' if edit_result.ok else '❌'} {edit_result.backend_used}")
        
        # Test guard tool
        guard_rules = [
            GuardRule(
                id="no-net-http-in-api",
                type="import",
                glob="api/*.go",
                forbid_import_prefix="net/http"
            )
        ]
        guard_result = await tools.guard(target_spec, GuardArgs(rules=guard_rules))
        print(f"  guard result: {'✅' if guard_result.ok else '❌'} violations: {len(guard_result.errors)}")
        
        # Test specific package
        pkg_target = TargetSpec(
            target="pkg:./cli/...",
            root=str(workspace_dir)
        )
        
        test_result = await tools.test(pkg_target, TestArgs(opts={"dry_run": True}))
        print(f"  test result: {'✅' if test_result.ok else '❌'} {test_result.stdout}")


async def test_typescript_project():
    """Test TypeScript project scenario"""
    print("🔧 Testing TypeScript project scenario...")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        workspace_dir = Path(temp_dir)
        
        # Create package.json with workspaces
        package_json = workspace_dir / "package.json"
        package_json.write_text(json.dumps({
            "name": "test-workspace",
            "workspaces": ["packages/*"],
            "scripts": {
                "test": "jest",
                "build": "tsc"
            },
            "devDependencies": {
                "typescript": "^5.0.0",
                "jest": "^29.0.0"
            }
        }, indent=2))
        
        # Create TypeScript config
        tsconfig = workspace_dir / "tsconfig.json"
        tsconfig.write_text(json.dumps({
            "compilerOptions": {
                "target": "ES2022",
                "module": "commonjs",
                "strict": True
            }
        }, indent=2))
        
        # Create packages
        packages_dir = workspace_dir / "packages"
        packages_dir.mkdir()
        
        # Core package
        core_dir = packages_dir / "core"
        core_dir.mkdir()
        
        core_package = core_dir / "package.json"
        core_package.write_text(json.dumps({
            "name": "@test/core",
            "version": "1.0.0"
        }, indent=2))
        
        core_index = core_dir / "index.ts"
        core_index.write_text("""export interface User {
  id: string;
  name: string;
}

export class UserService {
  getUser(id: string): User | null {
    return { id, name: 'Test User' };
  }
}
""")
        
        # Test TypeScript workspace
        target_spec = TargetSpec(
            target="ws",
            root=str(workspace_dir),
            language="ts"
        )
        
        fmt_result = await tools.fmt(target_spec, FmtArgs())
        print(f"  fmt result: {'✅' if fmt_result.ok else '❌'} {fmt_result.backend_used}")
        
        build_result = await tools.build(target_spec, BuildArgs(opts={"dry_run": True}))
        print(f"  build result: {'✅' if build_result.ok else '❌'} {build_result.stdout}")
        
        lint_result = await tools.lint(target_spec, LintArgs(opts={"dry_run": True}))
        print(f"  lint result: {'✅' if lint_result.ok else '❌'} {lint_result.backend_used}")


async def test_target_resolution():
    """Test target resolution scenarios"""
    print("🔧 Testing target resolution...")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        workspace_dir = Path(temp_dir)
        
        # Create simple Go project
        go_mod = workspace_dir / "go.mod"
        go_mod.write_text("module test\n\ngo 1.21\n")
        
        main_go = workspace_dir / "main.go"
        main_go.write_text("""package main

import "fmt"

func main() {
    fmt.Println("Hello")
}
""")
        
        src_dir = workspace_dir / "src"
        src_dir.mkdir()
        
        helper_go = src_dir / "helper.go"
        helper_go.write_text("""package src

func Helper() string {
    return "helper"
}
""")
        
        # Test different target types
        test_cases = [
            ("file:" + str(main_go), "Single file"),
            ("dir:" + str(src_dir), "Directory"),
            ("ws", "Workspace"),
            ("pkg:.", "Current package")
        ]
        
        for target, description in test_cases:
            target_spec = TargetSpec(
                target=target,
                root=str(workspace_dir),
                dry_run=True
            )
            
            result = await tools.fmt(target_spec, FmtArgs())
            print(f"  {description}: {'✅' if result.ok else '❌'} scope={len(result.scope_resolved)}")


async def test_composition_workflow():
    """Test tool composition workflow"""
    print("🔧 Testing composition workflow...")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        workspace_dir = Path(temp_dir)
        
        # Create Python project
        pyproject = workspace_dir / "pyproject.toml"
        pyproject.write_text("""[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[project]
name = "test-project"
version = "0.1.0"
""")
        
        src_dir = workspace_dir / "src"
        src_dir.mkdir()
        
        main_py = src_dir / "main.py"
        main_py.write_text("""import os
import sys
from typing import Optional


def process_user(user_id: str) -> Optional[str]:
    if not user_id:
        return None
    return f"User: {user_id}"


if __name__ == "__main__":
    print(process_user("123"))
""")
        
        # Composition workflow: edit -> fmt -> lint -> test
        target_spec = TargetSpec(
            target="file:" + str(main_py),
            root=str(workspace_dir),
            language="py"
        )
        
        # 1. Organize imports
        print("  Step 1: Organize imports")
        edit_result = await tools.edit(target_spec, EditArgs(op="organize_imports", dry_run=True))
        print(f"    {'✅' if edit_result.ok else '❌'} {edit_result.stdout}")
        
        # 2. Format code
        print("  Step 2: Format code")
        fmt_result = await tools.fmt(target_spec, FmtArgs())
        print(f"    {'✅' if fmt_result.ok else '❌'} {fmt_result.backend_used}")
        
        # 3. Lint code
        print("  Step 3: Lint code")
        lint_result = await tools.lint(target_spec, LintArgs(opts={"dry_run": True}))
        print(f"    {'✅' if lint_result.ok else '❌'} {lint_result.backend_used}")
        
        # 4. Test code
        print("  Step 4: Test code")
        test_result = await tools.test(target_spec, TestArgs(opts={"dry_run": True}))
        print(f"    {'✅' if test_result.ok else '❌'} {test_result.stdout}")


async def test_guard_rules():
    """Test guard rule scenarios"""
    print("🔧 Testing guard rules...")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        workspace_dir = Path(temp_dir)
        
        # Create SDK structure
        sdk_dir = workspace_dir / "sdk"
        sdk_dir.mkdir()
        
        # Bad import in SDK
        sdk_file = sdk_dir / "bad.py"
        sdk_file.write_text("""import node
from node.crypto import hash

def bad_function():
    return node.process()
""")
        
        # API contracts
        api_dir = workspace_dir / "api"
        api_dir.mkdir()
        
        api_file = api_dir / "contract.py"
        api_file.write_text("""import requests
from http.client import HTTPConnection

def api_call():
    return requests.get("http://example.com")
""")
        
        # Generated files
        generated_dir = workspace_dir / "api" / "pb"
        generated_dir.mkdir(parents=True)
        
        generated_file = generated_dir / "user_pb2.py"
        generated_file.write_text("# Generated file - do not edit")
        
        # Test guard rules
        guard_rules = [
            GuardRule(
                id="no-node-in-sdk",
                type="import",
                glob="sdk/**/*.py",
                forbid_import_prefix="node"
            ),
            GuardRule(
                id="no-http-in-contracts", 
                type="import",
                glob="api/*.py",
                forbid_import_prefix="requests"
            ),
            GuardRule(
                id="no-edits-in-generated",
                type="generated",
                glob="api/pb/**",
                forbid_writes=True
            )
        ]
        
        target_spec = TargetSpec(
            target="ws",
            root=str(workspace_dir)
        )
        
        guard_result = await tools.guard(target_spec, GuardArgs(rules=guard_rules))
        print(f"  guard result: {'✅' if guard_result.ok else '❌'}")
        print(f"  violations found: {len(guard_result.errors)}")
        for error in guard_result.errors:
            print(f"    - {error}")


async def main():
    """Run all tests"""
    print("🚀 Testing Exact 6-Tool Implementation\n")
    
    try:
        await test_go_workspace_scenario()
        print()
        
        await test_typescript_project() 
        print()
        
        await test_target_resolution()
        print()
        
        await test_composition_workflow()
        print()
        
        await test_guard_rules()
        print()
        
        print("🎉 All tests completed!")
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        raise


if __name__ == "__main__":
    asyncio.run(main())