Skip to content

Commit da8a1f0

Browse files
Saranya-jenaHarness
authored andcommitted
chore: [ML-1371]: Added get_audit_yaml tool and e2e tests (#159)
* fix: [ML-1371]: fixed UT Signed-off-by: Saranya-jena <[email protected]> * Merge branch 'master' of https://git0.harness.io/l7B_kbSEQD2wjrM7PShm5w/PROD/Harness_Commons/mcp-server into ML-1371 * chore: [ML-1371]: Added get_audit_yaml tool and e2e tests Signed-off-by: Saranya-jena <[email protected]>
1 parent 3532ed5 commit da8a1f0

File tree

6 files changed

+223
-6
lines changed

6 files changed

+223
-6
lines changed

client/audit.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,30 @@ import (
1010

1111
const (
1212
auditPath = "/api/audits/list"
13+
auditYamlPath = "/api/auditYaml"
1314
)
1415

1516
type AuditService struct {
1617
Client *Client
1718
}
1819

1920
// ListUserAuditTrail fetches the audit trail.
21+
// GetAuditYaml fetches the YAML diff for a specific audit event
22+
func (a *AuditService) GetAuditYaml(ctx context.Context, scope dto.Scope, auditID string) (*dto.AuditYamlResponse, error) {
23+
params := make(map[string]string)
24+
params["accountIdentifier"] = scope.AccountID
25+
params["routingId"] = scope.AccountID
26+
params["auditId"] = auditID
27+
28+
resp := &dto.AuditYamlResponse{}
29+
err := a.Client.Get(ctx, auditYamlPath, params, map[string]string{}, resp)
30+
if err != nil {
31+
return nil, fmt.Errorf("failed to get audit YAML: %w", err)
32+
}
33+
34+
return resp, nil
35+
}
36+
2037
func (a *AuditService) ListUserAuditTrail(ctx context.Context, scope dto.Scope, userIDList string, actionsList string, page int, size int, startTime int64, endTime int64, opts *dto.ListAuditEventsFilter) (*dto.AuditOutput[dto.AuditListItem], error) {
2138
if opts == nil {
2239
opts = &dto.ListAuditEventsFilter{}

client/dto/audit_yaml.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package dto
2+
3+
// AuditYamlResponse represents the response from the audit YAML API
4+
type AuditYamlResponse struct {
5+
Status string `json:"status"`
6+
Data AuditYamlData `json:"data"`
7+
MetaData interface{} `json:"metaData"`
8+
CorrelationID string `json:"correlationId"`
9+
}
10+
11+
// AuditYamlData contains the old and new YAML content for an audit event
12+
type AuditYamlData struct {
13+
OldYaml string `json:"oldYaml"`
14+
NewYaml string `json:"newYaml"`
15+
}

pkg/harness/tools/audit.go

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,45 @@ func previousWeek() string {
5757
return oneWeekAgo
5858
}
5959

60+
// GetAuditYamlTool creates a tool for retrieving YAML diff for a specific audit event.
61+
func GetAuditYamlTool(config *config.Config, auditClient *client.AuditService) (tool mcp.Tool, handler server.ToolHandlerFunc) {
62+
return mcp.NewTool("get_audit_yaml",
63+
mcp.WithDescription("Get YAML diff for a specific audit event."),
64+
mcp.WithString("audit_id",
65+
mcp.Description("The ID of the audit event to retrieve YAML diff for."),
66+
mcp.Required(),
67+
),
68+
WithScope(config, false),
69+
),
70+
func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
71+
slog.Info("Handling get_audit_yaml request", "request", request.GetArguments())
72+
73+
auditID, err := RequiredParam[string](request, "audit_id")
74+
if err != nil {
75+
return mcp.NewToolResultError(err.Error()), nil
76+
}
77+
78+
scope, err := FetchScope(config, request, false)
79+
if err != nil {
80+
return mcp.NewToolResultError(err.Error()), nil
81+
}
82+
83+
slog.Info("Calling GetAuditYaml API", "audit_id", auditID)
84+
85+
data, err := auditClient.GetAuditYaml(ctx, scope, auditID)
86+
if err != nil {
87+
return nil, fmt.Errorf("failed to get audit YAML: %w", err)
88+
}
89+
90+
r, err := json.Marshal(data)
91+
if err != nil {
92+
return nil, fmt.Errorf("failed to marshal the audit YAML response: %w", err)
93+
}
94+
95+
return mcp.NewToolResultText(string(r)), nil
96+
}
97+
}
98+
6099
// ListAuditsOfUser creates a tool for listing the audit trail.
61100
func ListUserAuditTrailTool(config *config.Config, auditClient *client.AuditService) (tool mcp.Tool, handler server.ToolHandlerFunc) {
62101
return mcp.NewTool("list_user_audits",

pkg/modules/core.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,7 @@ func RegisterAudit(config *config.Config, tsg *toolsets.ToolsetGroup) error {
253253
audit := toolsets.NewToolset("audit", "Audit log related tools").
254254
AddReadTools(
255255
toolsets.NewServerTool(tools.ListUserAuditTrailTool(config, auditService)),
256+
toolsets.NewServerTool(tools.GetAuditYamlTool(config, auditService)),
256257
)
257258

258259
// Add toolset to the group

test/e2e/audit_yaml_test.go

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
//go:build e2e
2+
3+
package e2e
4+
5+
import (
6+
"context"
7+
"encoding/json"
8+
"fmt"
9+
"testing"
10+
11+
"github.com/harness/harness-mcp/client/dto"
12+
"github.com/mark3labs/mcp-go/mcp"
13+
"github.com/stretchr/testify/require"
14+
)
15+
16+
// TestGetAuditYamlTool verifies that the get_audit_yaml tool is available
17+
func TestGetAuditYamlTool(t *testing.T) {
18+
t.Parallel()
19+
20+
mcpClient := setupMCPClient(t)
21+
ctx := context.Background()
22+
23+
// List available tools
24+
request := mcp.ListToolsRequest{}
25+
response, err := mcpClient.ListTools(ctx, request)
26+
require.NoError(t, err, "expected to list tools successfully")
27+
28+
// Check that audit yaml tool is available
29+
var foundAuditYamlTool bool
30+
31+
// Print all available tools
32+
fmt.Println("Available tools in TestGetAuditYamlTool:")
33+
for _, tool := range response.Tools {
34+
fmt.Printf("- %s\n", tool.Name)
35+
36+
// Check if this is our get_audit_yaml tool
37+
if tool.Name == "get_audit_yaml" {
38+
foundAuditYamlTool = true
39+
}
40+
}
41+
42+
// Check if we found the get_audit_yaml tool
43+
require.True(t, foundAuditYamlTool, "expected to find get_audit_yaml tool")
44+
}
45+
46+
// TestGetAuditYaml verifies that the get_audit_yaml tool works correctly
47+
func TestGetAuditYaml(t *testing.T) {
48+
t.Skip("Skipping TestGetAuditYaml as it requires specific audit events with YAML content")
49+
t.Parallel()
50+
51+
mcpClient := setupMCPClient(t)
52+
ctx := context.Background()
53+
accountID := getE2EAccountID(t)
54+
55+
// First, get an audit ID by listing audit events
56+
listRequest := mcp.CallToolRequest{}
57+
listRequest.Params.Name = "list_user_audits"
58+
listRequest.Params.Arguments = map[string]any{
59+
"accountIdentifier": accountID,
60+
"orgIdentifier": getE2EOrgID(),
61+
"projectIdentifier": getE2EProjectID(),
62+
"page": 0,
63+
"size": 1,
64+
"resource_type": "PIPELINE", // Filter for pipeline events which are likely to have YAML
65+
}
66+
67+
listResponse, err := mcpClient.CallTool(ctx, listRequest)
68+
require.NoError(t, err, "expected to call 'list_user_audits' tool successfully")
69+
if listResponse.IsError {
70+
t.Logf("Error response: %v", listResponse.Content)
71+
t.Log("list_user_audits tool returned an error")
72+
t.FailNow()
73+
}
74+
75+
// Parse the response to extract audit events
76+
if len(listResponse.Content) == 0 {
77+
t.Skip("No content returned from list_user_audits")
78+
}
79+
80+
textContent, ok := listResponse.Content[0].(mcp.TextContent)
81+
if !ok {
82+
t.Skip("Content is not of type TextContent")
83+
}
84+
85+
var auditResponse dto.AuditOutput[dto.AuditListItem]
86+
err = json.Unmarshal([]byte(textContent.Text), &auditResponse)
87+
if err != nil {
88+
t.Logf("Failed to unmarshal response: %v", err)
89+
t.Skip("Could not unmarshal audit response")
90+
}
91+
92+
// Skip test if no audit events found
93+
if len(auditResponse.Data.Content) == 0 {
94+
t.Skip("No audit events found to test get_audit_yaml")
95+
}
96+
97+
// Get the first audit ID
98+
auditID := auditResponse.Data.Content[0].AuditID
99+
if auditID == "" {
100+
t.Skip("Audit ID is empty")
101+
}
102+
t.Logf("Using audit ID: %s", auditID)
103+
104+
// Now call the get_audit_yaml tool with this audit ID
105+
yamlRequest := mcp.CallToolRequest{}
106+
yamlRequest.Params.Name = "get_audit_yaml"
107+
yamlRequest.Params.Arguments = map[string]any{
108+
"accountIdentifier": accountID,
109+
"orgIdentifier": getE2EOrgID(),
110+
"projectIdentifier": getE2EProjectID(),
111+
"audit_id": auditID,
112+
}
113+
114+
yamlResponse, err := mcpClient.CallTool(ctx, yamlRequest)
115+
require.NoError(t, err, "expected to call 'get_audit_yaml' tool successfully")
116+
117+
// Note: Some audit events might not have YAML content, so we don't strictly require success
118+
if yamlResponse.IsError {
119+
t.Logf("Error response from get_audit_yaml: %v", yamlResponse.Content)
120+
t.Skip("Skipping YAML content verification as the audit event might not have YAML content")
121+
}
122+
123+
// If we got a successful response, verify the structure
124+
require.NotEmpty(t, yamlResponse.Content, "expected content to not be empty")
125+
126+
// Parse the response to extract YAML content
127+
yamlTextContent, ok := yamlResponse.Content[0].(mcp.TextContent)
128+
require.True(t, ok, "expected content to be of type TextContent")
129+
130+
var auditYamlResponse dto.AuditYamlResponse
131+
err = json.Unmarshal([]byte(yamlTextContent.Text), &auditYamlResponse)
132+
require.NoError(t, err, "expected to unmarshal response successfully")
133+
134+
// Verify the response structure
135+
require.Equal(t, "SUCCESS", auditYamlResponse.Status, "expected status to be SUCCESS")
136+
137+
// Log the YAML content lengths for debugging
138+
oldYamlLen := len(auditYamlResponse.Data.OldYaml)
139+
newYamlLen := len(auditYamlResponse.Data.NewYaml)
140+
t.Logf("Old YAML length: %d, New YAML length: %d", oldYamlLen, newYamlLen)
141+
142+
// At least one of the YAMLs should be non-empty
143+
require.True(t, oldYamlLen > 0 || newYamlLen > 0,
144+
"expected at least one of old or new YAML to be non-empty")
145+
}

test/unit/register_prompts_test.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -111,8 +111,8 @@ func mockGetModulePrompts(fs embed.FS, module string, isInternal bool, mode stri
111111
return []prompts.PromptFile{
112112
{
113113
Metadata: prompts.PromptMetadata{
114-
Description: "Standard test description",
115-
ResultDescription: "Standard result description",
114+
Description: "Testmodule prompt description",
115+
ResultDescription: "Testmodule prompt result description",
116116
Module: "testmodule",
117117
},
118118
Content: "Standard content",
@@ -123,8 +123,8 @@ func mockGetModulePrompts(fs embed.FS, module string, isInternal bool, mode stri
123123
return []prompts.PromptFile{
124124
{
125125
Metadata: prompts.PromptMetadata{
126-
Description: "Architect test description",
127-
ResultDescription: "Architect result description",
126+
Description: "Testmodule prompt description",
127+
ResultDescription: "Testmodule prompt result description",
128128
Module: "testmodule",
129129
},
130130
Content: "Architect content",
@@ -155,10 +155,10 @@ func TestRegisterPrompts(t *testing.T) {
155155
if prompt.Name != "TESTMODULE" {
156156
t.Errorf("expected prompt name 'TESTMODULE', got %q", prompt.Name)
157157
}
158-
if !strings.Contains(prompt.Description, "Standard test description") {
158+
if !strings.Contains(prompt.Description, "Testmodule prompt description") {
159159
t.Errorf("description mismatch, got %q", prompt.Description)
160160
}
161-
if !strings.Contains(prompt.ResultDescription, "Standard result description") {
161+
if !strings.Contains(prompt.ResultDescription, "Testmodule prompt result description") {
162162
t.Errorf("result description mismatch, got %q", prompt.ResultDescription)
163163
}
164164

0 commit comments

Comments
 (0)