tests/test_tools_vault.py

"""Test zero-limitation tools, code editing, and storage vault."""

import sys
import os
import tempfile
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from splitbit_llm.harness.tools import (
    ToolRegistry, get_default_tools, tool_loop, parse_tool_calls,
    _tool_shell_exec, _tool_write_file, _tool_read_file,
    _tool_code_edit, _tool_list_dir, _tool_make_dir, _tool_delete_file,
)
from splitbit_llm.storage.vault import StorageVault


def test_full_terminal_control():
    """Test that shell_exec has zero limitations — can run any command."""
    # Test a basic command
    result = _tool_shell_exec("echo hello world")
    assert "hello world" in result, f"Expected 'hello world' in result: {result}"
    print(f"  echo: {result.strip()}")

    # Test that no commands are blocked
    result2 = _tool_shell_exec("python --version")
    assert "Error: dangerous command blocked" not in result2
    print(f"  python --version: {result2.strip()}")

    # Test pip command (would be blocked by old safety filter)
    result3 = _tool_shell_exec("echo pip install numpy")
    assert "dangerous" not in result3.lower()
    print(f"  pip command: not blocked")


def test_code_edit_tool():
    """Test code_edit tool — can modify files including its own framework."""
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
        f.write("value = 'old'\nprint(value)\n")
        path = f.name

    try:
        # Read original
        content = _tool_read_file(path)
        assert "old" in content

        # Edit the file
        result = _tool_code_edit(f'"{path}" "old" "new"')
        assert "Replaced" in result, f"Expected replacement: {result}"
        print(f"  Edit result: {result}")

        # Verify change
        content = _tool_read_file(path)
        assert "new" in content
        assert "old" not in content
        print(f"  Verified: old → new")

    finally:
        os.unlink(path)


def test_list_dir_tool():
    """Test list_dir tool."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Create some files
        open(os.path.join(tmpdir, "file1.txt"), "w").close()
        open(os.path.join(tmpdir, "file2.py"), "w").close()
        os.makedirs(os.path.join(tmpdir, "subdir"))

        result = _tool_list_dir(tmpdir)
        assert "file1.txt" in result
        assert "file2.py" in result
        assert "subdir/" in result
        print(f"  Listed: {result.strip()}")


def test_make_dir_tool():
    """Test make_dir tool."""
    with tempfile.TemporaryDirectory() as tmpdir:
        new_dir = os.path.join(tmpdir, "new_project", "src")
        result = _tool_make_dir(new_dir)
        assert "Created" in result
        assert os.path.exists(new_dir)
        print(f"  Created: {new_dir}")


def test_delete_file_tool():
    """Test delete_file tool."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Test file deletion
        file_path = os.path.join(tmpdir, "temp.txt")
        open(file_path, "w").close()
        assert os.path.exists(file_path)
        result = _tool_delete_file(file_path)
        assert "Deleted file" in result
        assert not os.path.exists(file_path)

        # Test directory deletion
        dir_path = os.path.join(tmpdir, "temp_dir")
        os.makedirs(dir_path)
        result2 = _tool_delete_file(dir_path)
        assert "Deleted directory" in result2
        assert not os.path.exists(dir_path)
        print(f"  Deleted file and directory: OK")


def test_zero_limitation_tools_count():
    """Test that all zero-limitation tools are registered."""
    tools = get_default_tools()
    tool_names = [t.name for t in tools]
    assert "shell_exec" in tool_names
    assert "write_file" in tool_names
    assert "code_edit" in tool_names
    assert "list_dir" in tool_names
    assert "make_dir" in tool_names
    assert "delete_file" in tool_names
    assert "read_file" in tool_names
    assert "calculate" in tool_names
    print(f"  Tools: {len(tools)} registered — {tool_names}")

    # Check shell_exec description says full control
    shell_tool = next(t for t in tools if t.name == "shell_exec")
    assert "full control" in shell_tool.description.lower() or "no restriction" in shell_tool.description.lower()
    print(f"  shell_exec: {shell_tool.description}")


def test_storage_vault():
    """Test mass storage vault — auto-resizing storage."""
    with tempfile.TemporaryDirectory() as tmpdir:
        vault = StorageVault(data_dir=tmpdir)

        # Check directories created
        assert os.path.exists(vault.db_dir)
        assert os.path.exists(vault.cache_dir)
        assert os.path.exists(vault.archive_dir)
        assert os.path.exists(vault.artifacts_dir)
        print(f"  Directories: db, cache, archive, artifacts — all created")

        # Store and load artifact
        path = vault.store_artifact("test_code.py", b"print('hello')")
        assert os.path.exists(path)
        loaded = vault.load_artifact("test_code.py")
        assert loaded == b"print('hello')"
        print(f"  Artifact stored and loaded: {loaded}")

        # Store text artifact
        path2 = vault.store_artifact_text("notes.txt", "Important notes")
        loaded2 = vault.load_artifact("notes.txt")
        assert loaded2 == b"Important notes"
        print(f"  Text artifact: {loaded2}")

        # List artifacts
        artifacts = vault.list_artifacts()
        assert len(artifacts) == 2
        names = [a["name"] for a in artifacts]
        assert "test_code.py" in names
        assert "notes.txt" in names
        print(f"  Artifacts listed: {names}")

        # Delete artifact
        deleted = vault.delete_artifact("test_code.py")
        assert deleted is True
        artifacts = vault.list_artifacts()
        assert len(artifacts) == 1
        print(f"  Deleted: test_code.py ({len(artifacts)} remaining)")

        # Stats
        stats = vault.get_stats()
        assert stats["disk_total_gb"] > 0
        assert "total_storage_mb" in stats
        assert "disk_usage_percent" in stats
        print(f"  Stats: {stats['total_storage_mb']}MB used, "
              f"{stats['disk_free_gb']}GB free, {stats['disk_usage_percent']:.1%} disk usage")


def test_vault_auto_resize():
    """Test vault auto-resize and cleanup."""
    with tempfile.TemporaryDirectory() as tmpdir:
        vault = StorageVault(data_dir=tmpdir)

        # Store many artifacts
        for i in range(20):
            vault.store_artifact_text(f"artifact_{i}.txt", f"content {i}" * 100)

        stats = vault.get_stats()
        assert stats["artifacts_storage_bytes"] > 0
        print(f"  Stored 20 artifacts: {stats['artifacts_storage_bytes']} bytes")

        # Force cleanup (should not delete recent files)
        result = vault.force_cleanup()
        assert "cleanups_performed" in result
        # Recent files should still be there
        artifacts = vault.list_artifacts()
        assert len(artifacts) == 20  # nothing deleted (all recent)
        print(f"  Cleanup: {result['cleanups_performed']} cycles, {len(artifacts)} artifacts remain")


def test_vault_db_paths():
    """Test vault DB path management."""
    with tempfile.TemporaryDirectory() as tmpdir:
        vault = StorageVault(data_dir=tmpdir)

        db_path = vault.get_db_path("memory")
        assert db_path.endswith("memory.db")
        assert "db" in db_path

        cache_path = vault.get_cache_path("test.cache")
        assert "cache" in cache_path

        artifact_path = vault.get_artifact_path("code.py")
        assert "artifacts" in artifact_path
        print(f"  DB path: {db_path}")
        print(f"  Cache path: {cache_path}")
        print(f"  Artifact path: {artifact_path}")


def test_tool_loop_with_new_tools():
    """Test that tool_loop works with the new tools."""
    registry = ToolRegistry()
    for tool in get_default_tools():
        registry.register(tool)

    # Test list_dir via tool loop
    text = '[TOOL: list_dir(".")]'
    result_text, results = tool_loop(text, registry)
    assert len(results) > 0
    assert results[0].success
    print(f"  Tool loop: {results[0].name} → {results[0].output[:50]}...")


if __name__ == "__main__":
    print("Running zero-limitation tools and storage vault tests...")
    test_full_terminal_control()
    print("  ✓ test_full_terminal_control")
    test_code_edit_tool()
    print("  ✓ test_code_edit_tool")
    test_list_dir_tool()
    print("  ✓ test_list_dir_tool")
    test_make_dir_tool()
    print("  ✓ test_make_dir_tool")
    test_delete_file_tool()
    print("  ✓ test_delete_file_tool")
    test_zero_limitation_tools_count()
    print("  ✓ test_zero_limitation_tools_count")
    test_storage_vault()
    print("  ✓ test_storage_vault")
    test_vault_auto_resize()
    print("  ✓ test_vault_auto_resize")
    test_vault_db_paths()
    print("  ✓ test_vault_db_paths")
    test_tool_loop_with_new_tools()
    print("  ✓ test_tool_loop_with_new_tools")
    print("\nAll zero-limitation & vault tests passed!")