codemode/codemode_test.go at main · Protocol-Lattice/codemode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package codemode_test

import (
	"context"
	"fmt"
	"testing"

	"github.com/Protocol-Lattice/codemode"
	"github.com/mark3labs/mcp-go/mcp"
)

type mockMCPClient struct {
	tools map[string]func(args map[string]any) (*mcp.CallToolResult, error)
}

func (m *mockMCPClient) ListTools(ctx context.Context, req mcp.ListToolsRequest) (*mcp.ListToolsResult, error) {
	return &mcp.ListToolsResult{Tools: []mcp.Tool{}}, nil
}

func (m *mockMCPClient) CallTool(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	handler, ok := m.tools[req.Params.Name]
	if !ok {
		return nil, fmt.Errorf("tool not found: %s", req.Params.Name)
	}
	args, _ := req.Params.Arguments.(map[string]any)
	return handler(args)
}

func TestCodeModeExecute_Basic(t *testing.T) {
	cm := codemode.NewCodeModeMCP(&mockMCPClient{})

	res, err := cm.Execute(context.Background(), codemode.CodeModeArgs{
		Code:    "__out = 1 + 1",
		Timeout: 5000,
	})

	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if res.Value != 2 {
		t.Errorf("expected 2, got %v", res.Value)
	}
}

func TestCodeModeExecute_CallTool(t *testing.T) {
	client := &mockMCPClient{
		tools: map[string]func(args map[string]any) (*mcp.CallToolResult, error){
			"echo": func(args map[string]any) (*mcp.CallToolResult, error) {
				msg := args["msg"].(string)
				return mcp.NewToolResultText("echoed: " + msg), nil
			},
		},
	}
	cm := codemode.NewCodeModeMCP(client)

	code := `
		res, err := codemode.CallTool("echo", map[string]any{"msg": "hello"})
		if err != nil {
			panic(err)
		}
		__out = res
	`

	res, err := cm.Execute(context.Background(), codemode.CodeModeArgs{
		Code:    code,
		Timeout: 5000,
	})

	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if res.Value != "echoed: hello" {
		t.Errorf("expected 'echoed: hello', got %v", res.Value)
	}
}

func TestCodeModeExecute_Timeout(t *testing.T) {
	cm := codemode.NewCodeModeMCP(&mockMCPClient{})

	code := `
		for {
			// infinite loop
		}
	`

	_, err := cm.Execute(context.Background(), codemode.CodeModeArgs{
		Code:    code,
		Timeout: 50,
	})

	if err == nil {
		t.Fatal("expected timeout error, got nil")
	}

	if err.Error() != "execution timed out after 50ms" {
		t.Errorf("unexpected error message: %v", err)
	}
}

func TestCodeModeExecute_Panic(t *testing.T) {
	cm := codemode.NewCodeModeMCP(&mockMCPClient{})

	code := `
		panic("test panic")
	`

	_, err := cm.Execute(context.Background(), codemode.CodeModeArgs{
		Code:    code,
		Timeout: 5000,
	})

	if err == nil {
		t.Fatal("expected panic error, got nil")
	}
}