66 lines
1.7 KiB
Go
66 lines
1.7 KiB
Go
package pii
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/agent"
|
|
"github.com/sipeed/picoclaw/pkg/providers"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestRedactor_Redact(t *testing.T) {
|
|
r := NewRedactor(true)
|
|
|
|
tests := []struct {
|
|
input string
|
|
expected string
|
|
}{
|
|
{"Hello, contact me at steve@example.com", "Hello, contact me at [EMAIL_1]"},
|
|
{"My IP is 192.168.1.1", "My IP is [IP_1]"},
|
|
{"Call me at +1 555-123-4567", "Call me at [PHONE_1]"},
|
|
{"Nothing sensitive here", "Nothing sensitive here"},
|
|
}
|
|
|
|
mapping := r.getMapping("test")
|
|
for _, tt := range tests {
|
|
assert.Equal(t, tt.expected, r.redact(tt.input, mapping))
|
|
}
|
|
}
|
|
|
|
func TestRedactor_BeforeLLM(t *testing.T) {
|
|
r := NewRedactor(true)
|
|
ctx := context.Background()
|
|
|
|
req := &agent.LLMHookRequest{
|
|
Messages: []providers.Message{
|
|
{Role: "user", Content: "My email is user@foo.com"},
|
|
{Role: "system", Content: "Keep 127.0.0.1"}, // system message should not be redacted
|
|
},
|
|
}
|
|
|
|
next, decision, err := r.BeforeLLM(ctx, req)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, agent.HookActionContinue, decision.Action)
|
|
|
|
assert.Equal(t, "My email is [EMAIL_1]", next.Messages[0].Content)
|
|
assert.Equal(t, "Keep 127.0.0.1", next.Messages[1].Content)
|
|
}
|
|
|
|
func TestRedactor_AfterLLM(t *testing.T) {
|
|
r := NewRedactor(true)
|
|
ctx := context.Background()
|
|
|
|
resp := &agent.LLMHookResponse{
|
|
Response: &providers.LLMResponse{
|
|
Content: "The user's email was user@foo.com",
|
|
},
|
|
}
|
|
|
|
next, decision, err := r.AfterLLM(ctx, resp)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, agent.HookActionContinue, decision.Action)
|
|
|
|
assert.Equal(t, "The user's email was user@foo.com", next.Response.Content)
|
|
}
|