[AI] Expand Test Coverage - tavern/internal/mcp

tavern/internal/mcp/export_test.go

@@ -0,0 +1,12 @@
+package mcp
+
+// Export unexported types and functions for testing
+
+type ContextKey = contextKey
+
+var (
+	HandleListHosts   = handleListHosts
+	HandleListQuests  = handleListQuests
+	HandleListTomes   = handleListTomes
+	HandleQuestOutput = handleQuestOutput
+)

tavern/internal/mcp/tool_list_hosts_test.go

@@ -0,0 +1,116 @@
+package mcp_test
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"realm.pub/tavern/internal/c2/c2pb"
+	"realm.pub/tavern/internal/ent/tag"
+	tavernmcp "realm.pub/tavern/internal/mcp"
+)
+
+func TestHandleListHosts(t *testing.T) {
+	client := setupTestDB(t)
+	defer client.Close()
+
+	ctx := context.WithValue(context.Background(), tavernmcp.ContextKey{}, client)
+
+	// Create tags
+	tag1 := client.Tag.Create().SetName("web").SetKind(tag.KindService).SaveX(ctx)
+	tag2 := client.Tag.Create().SetName("db").SetKind(tag.KindService).SaveX(ctx)
+
+	// Create mock hosts
+	host1 := client.Host.Create().
+		SetIdentifier("host-1").
+		SetPlatform(c2pb.Host_PLATFORM_LINUX).
+		SetName("web-server").
+		SetPrimaryIP("192.168.1.10").
+		SetExternalIP("203.0.113.10").
+		SetLastSeenAt(time.Now()).
+		AddTags(tag1).
+		SaveX(ctx)
+
+	// Create mock beacons and attach them to host1
+	client.Beacon.Create().
+		SetHost(host1).
+		SetName("beacon-1").
+		SetTransport(c2pb.Transport_TRANSPORT_HTTP1).
+		SetPrincipal("admin").
+		SaveX(ctx)
+
+	host2 := client.Host.Create().
+		SetIdentifier("host-2").
+		SetPlatform(c2pb.Host_PLATFORM_WINDOWS).
+		SetName("db-server").
+		SetPrimaryIP("192.168.1.20").
+		SetExternalIP("203.0.113.20").
+		SetLastSeenAt(time.Now()).
+		AddTags(tag2).
+		SaveX(ctx)
+	_ = host2
+
+	req := mcp.CallToolRequest{}
+	res, err := tavernmcp.HandleListHosts(ctx, req)
+	require.NoError(t, err)
+	require.NotNil(t, res)
+	require.False(t, res.IsError)
+
+	var results []map[string]interface{}
+	textContent := res.Content[0].(mcp.TextContent).Text
+	err = json.Unmarshal([]byte(textContent), &results)
+	require.NoError(t, err)
+
+	assert.Len(t, results, 2)
+
+	// Verify host 1
+	var h1Res map[string]interface{}
+	var h2Res map[string]interface{}
+	if results[0]["id"].(float64) == float64(host1.ID) {
+		h1Res = results[0]
+		h2Res = results[1]
+	} else {
+		h1Res = results[1]
+		h2Res = results[0]
+	}
+
+	assert.Equal(t, "host-1", h1Res["identifier"])
+	assert.Equal(t, "web-server", h1Res["name"])
+	assert.Equal(t, "PLATFORM_LINUX", h1Res["platform"])
+	assert.Equal(t, "192.168.1.10", h1Res["primaryIP"])
+	assert.Equal(t, "203.0.113.10", h1Res["externalIP"])
+
+	tags1 := h1Res["tags"].([]interface{})
+	require.Len(t, tags1, 1)
+	tagRes := tags1[0].(map[string]interface{})
+	assert.Equal(t, "web", tagRes["name"])
+	assert.Equal(t, "service", tagRes["kind"])
+
+	beacons1 := h1Res["beacons"].([]interface{})
+	require.Len(t, beacons1, 1)
+	beaconRes := beacons1[0].(map[string]interface{})
+	assert.Equal(t, "beacon-1", beaconRes["name"])
+	assert.Equal(t, "admin", beaconRes["principal"])
+
+	// Verify host 2
+	assert.Equal(t, "host-2", h2Res["identifier"])
+	assert.Equal(t, "db-server", h2Res["name"])
+	assert.Equal(t, "PLATFORM_WINDOWS", h2Res["platform"])
+	assert.Equal(t, "192.168.1.20", h2Res["primaryIP"])
+	assert.Equal(t, "203.0.113.20", h2Res["externalIP"])
+	tags2 := h2Res["tags"].([]interface{})
+	require.Len(t, tags2, 1)
+	tagRes2 := tags2[0].(map[string]interface{})
+	assert.Equal(t, "db", tagRes2["name"])
+	assert.Equal(t, "service", tagRes2["kind"])
+
+	// test no client error
+	res, err = tavernmcp.HandleListHosts(context.Background(), req)
+	require.NoError(t, err)
+	require.True(t, res.IsError)
+	assert.Contains(t, res.Content[0].(mcp.TextContent).Text, "internal error: no database client")
+}

tavern/internal/mcp/tool_list_quests_test.go

@@ -0,0 +1,105 @@
+package mcp_test
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"realm.pub/tavern/internal/ent/tome"
+	tavernmcp "realm.pub/tavern/internal/mcp"
+)
+
+func TestHandleListQuests(t *testing.T) {
+	client := setupTestDB(t)
+	defer client.Close()
+
+	ctx := context.WithValue(context.Background(), tavernmcp.ContextKey{}, client)
+
+	// Create user
+	user1 := client.User.Create().
+		SetOauthID("oauth-1").
+		SetName("Alice").
+		SetPhotoURL("http://example.com/photo").
+		SaveX(ctx)
+
+	// Create mock tomes
+	tome1 := client.Tome.Create().
+		SetName("discovery-tome").
+		SetTactic(tome.TacticDISCOVERY).
+		SetDescription("Discover things").
+		SetParamDefs(`[{"name": "p1", "type":"string"}]`).
+		SetAuthor("Alice").
+		SetEldritch("console.log('hello')").
+		SaveX(ctx)
+
+	tome2 := client.Tome.Create().
+		SetName("exec-tome").
+		SetTactic(tome.TacticEXECUTION).
+		SetDescription("Execute things").
+		SetParamDefs(`[]`).
+		SetAuthor("Bob").
+		SetEldritch("console.log('world')").
+		SaveX(ctx)
+
+	quest1 := client.Quest.Create().
+		SetName("test-quest-1").
+		SetParameters(`{"p1": "v1"}`).
+		SetTome(tome1).
+		SetCreator(user1).
+		SaveX(ctx)
+
+	quest2 := client.Quest.Create().
+		SetName("test-quest-2").
+		SetParameters(`{}`).
+		SetTome(tome2).
+		SaveX(ctx)
+	_ = quest2
+
+	req := mcp.CallToolRequest{}
+	res, err := tavernmcp.HandleListQuests(ctx, req)
+	require.NoError(t, err)
+	require.NotNil(t, res)
+	require.False(t, res.IsError)
+
+	var results []map[string]interface{}
+	textContent := res.Content[0].(mcp.TextContent).Text
+	err = json.Unmarshal([]byte(textContent), &results)
+	require.NoError(t, err)
+
+	assert.Len(t, results, 2)
+
+	var q1Res map[string]interface{}
+	var q2Res map[string]interface{}
+	if results[0]["id"].(float64) == float64(quest1.ID) {
+		q1Res = results[0]
+		q2Res = results[1]
+	} else {
+		q1Res = results[1]
+		q2Res = results[0]
+	}
+
+	assert.Equal(t, "test-quest-1", q1Res["name"])
+	assert.Equal(t, "Alice", q1Res["creator"])
+	assert.Equal(t, `{"p1": "v1"}`, q1Res["parameters"])
+	assert.Equal(t, "discovery-tome", q1Res["tomeName"])
+	assert.Equal(t, "DISCOVERY", q1Res["tomeTactic"])
+	assert.Equal(t, "Discover things", q1Res["tomeDescription"])
+	assert.NotEmpty(t, q1Res["createdAt"])
+
+	assert.Equal(t, "test-quest-2", q2Res["name"])
+	assert.NotContains(t, q2Res, "creator")
+	assert.Equal(t, `{}`, q2Res["parameters"])
+	assert.Equal(t, "exec-tome", q2Res["tomeName"])
+	assert.Equal(t, "EXECUTION", q2Res["tomeTactic"])
+	assert.Equal(t, "Execute things", q2Res["tomeDescription"])
+	assert.NotEmpty(t, q2Res["createdAt"])
+
+	// test no client error
+	res, err = tavernmcp.HandleListQuests(context.Background(), req)
+	require.NoError(t, err)
+	require.True(t, res.IsError)
+	assert.Contains(t, res.Content[0].(mcp.TextContent).Text, "internal error: no database client")
+}

tavern/internal/mcp/tool_list_tomes_test.go

@@ -0,0 +1,79 @@
+package mcp_test
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"realm.pub/tavern/internal/ent/tome"
+	tavernmcp "realm.pub/tavern/internal/mcp"
+)
+
+func TestHandleListTomes(t *testing.T) {
+	client := setupTestDB(t)
+	defer client.Close()
+
+	ctx := context.WithValue(context.Background(), tavernmcp.ContextKey{}, client)
+
+	// Create mock tomes
+	tome1 := client.Tome.Create().
+		SetName("discovery-tome").
+		SetTactic(tome.TacticDISCOVERY).
+		SetDescription("Discover things").
+		SetParamDefs(`[{"name": "p1", "type":"string"}]`).
+		SetAuthor("Alice").
+		SetEldritch("console.log('hello')").
+		SaveX(ctx)
+
+	tome2 := client.Tome.Create().
+		SetName("exec-tome").
+		SetTactic(tome.TacticEXECUTION).
+		SetDescription("Execute things").
+		SetParamDefs(`[]`).
+		SetAuthor("Bob").
+		SetEldritch("console.log('world')").
+		SaveX(ctx)
+	_ = tome2
+
+	req := mcp.CallToolRequest{}
+	res, err := tavernmcp.HandleListTomes(ctx, req)
+	require.NoError(t, err)
+	require.NotNil(t, res)
+	require.False(t, res.IsError)
+
+	var results []map[string]interface{}
+	textContent := res.Content[0].(mcp.TextContent).Text
+	err = json.Unmarshal([]byte(textContent), &results)
+	require.NoError(t, err)
+
+	assert.Len(t, results, 2)
+
+	var t1Res map[string]interface{}
+	var t2Res map[string]interface{}
+	if results[0]["id"].(float64) == float64(tome1.ID) {
+		t1Res = results[0]
+		t2Res = results[1]
+	} else {
+		t1Res = results[1]
+		t2Res = results[0]
+	}
+
+	assert.Equal(t, "discovery-tome", t1Res["name"])
+	assert.Equal(t, "DISCOVERY", t1Res["tactic"])
+	assert.Equal(t, "Discover things", t1Res["description"])
+	assert.Equal(t, `[{"name": "p1", "type":"string"}]`, t1Res["paramDefs"])
+
+	assert.Equal(t, "exec-tome", t2Res["name"])
+	assert.Equal(t, "EXECUTION", t2Res["tactic"])
+	assert.Equal(t, "Execute things", t2Res["description"])
+	assert.Equal(t, `[]`, t2Res["paramDefs"])
+
+	// test no client error
+	res, err = tavernmcp.HandleListTomes(context.Background(), req)
+	require.NoError(t, err)
+	require.True(t, res.IsError)
+	assert.Contains(t, res.Content[0].(mcp.TextContent).Text, "internal error: no database client")
+}