test_tool_registry.py

#!/usr/bin/env python3
"""
Test Tool Registry

Verifies that:
1. Default tools are registered
2. Builder is aware of tools
3. Generated code uses tools
"""
from dynograph.tools import default_tool_registry
from dynograph.generation.builder import GraphBuilder


def test_default_registry():
    """Test that default tool registry has all tools."""
    print("=" * 80)
    print("TEST 1: Default Tool Registry")
    print("=" * 80)

    print(f"\nRegistry: {default_tool_registry}")
    print(f"Total tools: {len(default_tool_registry)}")

    tools = default_tool_registry.list_tools()
    descriptions = default_tool_registry.get_tool_descriptions()

    print(f"\n Available Tools ({len(tools)}):")
    print("-" * 80)

    # Group by category
    categories = {
        "Search & Research": ["web_search", "wikipedia"],
        "Code & Execution": ["code", "python_repl"],
        "Math & Analysis": ["calculator", "text_analyzer"],
        "File Operations": ["file_reader", "file_writer"],
        "Network & API": ["http_request"],
        "System": ["shell", "datetime"],
        "Data Processing": ["json_parser"],
    }

    for category, tool_names in categories.items():
        print(f"\n{category}:")
        for tool_name in tool_names:
            if tool_name in descriptions:
                print(f"  ✓ {tool_name}: {descriptions[tool_name]}")
            else:
                print(f"  ✗ {tool_name}: NOT FOUND")

    # Verify expected count
    expected = 12
    if len(tools) == expected:
        print(f"\n✅ All {expected} default tools registered")
        return True
    else:
        print(f"\n❌ Expected {expected} tools, found {len(tools)}")
        return False


def test_builder_integration():
    """Test that GraphBuilder uses tool registry."""
    print("\n" + "=" * 80)
    print("TEST 2: Builder Integration")
    print("=" * 80)

    # Create builder (should auto-use default registry)
    builder = GraphBuilder()

    if builder.tool_registry is None:
        print("\n❌ Builder has no tool registry!")
        return False

    tools = builder.tool_registry.list_tools()
    print(f"\n✅ Builder has tool registry with {len(tools)} tools")

    # Verify it's the default registry
    if builder.tool_registry == default_tool_registry:
        print("✅ Using default_tool_registry")
    else:
        print("ℹ️  Using custom registry")

    return True


def test_tool_in_prompt():
    """Test that tools are shown in generation prompts."""
    print("\n" + "=" * 80)
    print("TEST 3: Tools in Prompts")
    print("=" * 80)

    from dynograph.generation.prompts import build_code_generation_prompt

    # Get tools from registry
    tools_dict = default_tool_registry.get_tool_descriptions()

    # Build a simple prompt
    prompt = build_code_generation_prompt(
        user_query="Create a graph that searches the web",
        examples=[],  # Empty for this test
        available_tools=tools_dict
    )

    # Check if tools are mentioned
    checks = {
        "web_search in prompt": "web_search" in prompt,
        "code in prompt": "code" in prompt,
        "file_reader in prompt": "file_reader" in prompt,
        "Tool count shown": f"{len(tools_dict)} tools" in prompt,
        "get_tool mentioned": "get_tool" in prompt,
    }

    print("\nPrompt checks:")
    all_pass = True
    for check, result in checks.items():
        status = "✅" if result else "❌"
        print(f"  {status} {check}")
        if not result:
            all_pass = False

    if all_pass:
        print("\n✅ Tools properly displayed in prompts")
    else:
        print("\n❌ Some tools missing from prompts")

    # Show snippet of tools section
    if "AVAILABLE TOOLS" in prompt:
        tools_start = prompt.find("AVAILABLE TOOLS")
        tools_section = prompt[tools_start:tools_start+500]
        print("\nTools section preview:")
        print("-" * 80)
        print(tools_section[:400] + "...")
        print("-" * 80)

    return all_pass


def test_tool_access():
    """Test that we can actually get and use tools."""
    print("\n" + "=" * 80)
    print("TEST 4: Tool Access")
    print("=" * 80)

    from dynograph.tools import get_tool

    # Test getting a few tools
    test_tools = ["text_analyzer", "datetime", "json_parser"]

    print("\nTesting tool retrieval:")
    all_success = True

    for tool_name in test_tools:
        tool = get_tool(tool_name)
        if tool:
            print(f"  ✅ {tool_name}: {type(tool).__name__}")

            # Try to use it
            try:
                if tool_name == "text_analyzer":
                    result = tool.invoke("Hello world test")
                    print(f"     → Result: {result}")
                elif tool_name == "datetime":
                    result = tool.invoke("%Y-%m-%d")
                    print(f"     → Result: {result}")
                elif tool_name == "json_parser":
                    result = tool.invoke('{"test": "value"}')
                    print(f"     → Result: {result}")
            except Exception as e:
                print(f"     ⚠️  Error using tool: {e}")
        else:
            print(f"  ❌ {tool_name}: Could not retrieve")
            all_success = False

    if all_success:
        print("\n✅ All tools accessible and working")
    else:
        print("\n❌ Some tools failed")

    return all_success


def main():
    """Run all tests."""
    print("\n" + "╔" + "=" * 78 + "╗")
    print("║" + "  TOOL REGISTRY TESTS".center(78) + "║")
    print("╚" + "=" * 78 + "╝")

    results = {
        "Default Registry": test_default_registry(),
        "Builder Integration": test_builder_integration(),
        "Tools in Prompts": test_tool_in_prompt(),
        "Tool Access": test_tool_access(),
    }

    # Summary
    print("\n" + "=" * 80)
    print("SUMMARY")
    print("=" * 80)

    for test, passed in results.items():
        status = "✅ PASS" if passed else "❌ FAIL"
        print(f"  {status}  {test}")

    passed_count = sum(1 for v in results.values() if v)
    total = len(results)

    print(f"\n{passed_count}/{total} tests passed")

    if passed_count == total:
        print("\n🎉 All tool registry tests passed!")
        return True
    else:
        print("\n⚠️  Some tests failed")
        return False


if __name__ == "__main__":
    import sys
    success = main()
    sys.exit(0 if success else 1)