@@ -1,34 +1,35 @@
11package mcp
22
33import (
4+ "bytes"
45 "encoding/base64"
56 "testing"
67
78 "github.com/stretchr/testify/require"
89)
910
10−func TestEnsureBase64(t *testing.T) {
11+func TestEnsureRawBytes(t *testing.T) {
1112 t.Parallel()
1213
1314 tests := []struct {
1415 name string
1516 input []byte
16− wantData []byte // expected output
17+ wantData []byte
1718 }{
1819 {
1920 name: "already base64 encoded",
2021 input: []byte("SGVsbG8gV29ybGQh"), // "Hello World!" in base64
21− wantData: []byte("SGVsbG8gV29ybGQh"),
22+ wantData: []byte("Hello World!"),
2223 },
2324 {
2425 name: "raw binary data (PNG header)",
2526 input: []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A},
26− wantData: []byte(base64.StdEncoding.EncodeToString([]byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A})),
27+ wantData: []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A},
2728 },
2829 {
2930 name: "raw binary with high bytes",
3031 input: []byte{0xFF, 0xD8, 0xFF, 0xE0}, // JPEG header
31− wantData: []byte(base64.StdEncoding.EncodeToString([]byte{0xFF, 0xD8, 0xFF, 0xE0})),
32+ wantData: []byte{0xFF, 0xD8, 0xFF, 0xE0},
3233 },
3334 {
3435 name: "empty data",
@@ -38,88 +39,31 @@ func TestEnsureBase64(t *testing.T) {
3839 {
3940 name: "base64 with padding",
4041 input: []byte("YQ=="), // "a" in base64
41− wantData: []byte("YQ=="),
42+ wantData: []byte("a"),
4243 },
4344 {
4445 name: "base64 without padding",
4546 input: []byte("YQ"),
46− wantData: []byte("YQ=="),
47+ wantData: []byte("a"),
4748 },
4849 {
4950 name: "base64 with whitespace",
5051 input: []byte("U0dWc2JHOGdWMjl5YkdRaA==\n"),
51− wantData: []byte("U0dWc2JHOGdWMjl5YkdRaA=="),
52+ wantData: []byte("SGVsbG8gV29ybGQh"),
5253 },
5354 }
5455
5556 for _, tt := range tests {
5657 t.Run(tt.name, func(t *testing.T) {
5758 t.Parallel()
58− result := ensureBase64(tt.input)
59+ result := ensureRawBytes(tt.input)
5960 require.Equal(t, tt.wantData, result)
6061
61− // Verify the result is valid base64 that can be decoded.
62− if len(result) > 0 {
63− _, err := base64.StdEncoding.DecodeString(string(result))
64− if err != nil {
65− _, err = base64.RawStdEncoding.DecodeString(string(result))
66− }
67− require.NoError(t, err, "result should be valid base64")
62+ if len(result) > 0 && !bytes.Equal(result, tt.input) {
63+ reEncoded := base64.StdEncoding.EncodeToString(result)
64+ _, err := base64.StdEncoding.DecodeString(reEncoded)
65+ require.NoError(t, err, "re-encoded result should be valid base64")
6866 }
6967 })
7068 }
7169}
72−
73−func TestIsValidBase64(t *testing.T) {
74− t.Parallel()
75−
76− tests := []struct {
77− name string
78− input []byte
79− want bool
80− }{
81− {
82− name: "valid base64",
83− input: []byte("SGVsbG8gV29ybGQh"),
84− want: true,
85− },
86− {
87− name: "valid base64 with padding",
88− input: []byte("YQ=="),
89− want: true,
90− },
91− {
92− name: "raw binary with high bytes",
93− input: []byte{0xFF, 0xD8, 0xFF},
94− want: false,
95− },
96− {
97− name: "empty",
98− input: []byte{},
99− want: true,
100− },
101− {
102− name: "valid raw base64 without padding",
103− input: []byte("YQ"),
104− want: true,
105− },
106− {
107− name: "valid base64 with whitespace",
108− input: normalizeBase64Input([]byte("U0dWc2JHOGdWMjl5YkdRaA==\n")),
109− want: true,
110− },
111− {
112− name: "invalid base64 characters",
113− input: []byte("SGVsbG8!@#$"),
114− want: false,
115− },
116− }
117−
118− for _, tt := range tests {
119− t.Run(tt.name, func(t *testing.T) {
120− t.Parallel()
121− got := isValidBase64(tt.input)
122− require.Equal(t, tt.want, got)
123− })
124− }
125−}