content_test.go

 1package message
 2
 3import (
 4	"fmt"
 5	"strings"
 6	"testing"
 7)
 8
 9func makeTestAttachments(n int, contentSize int) []Attachment {
10	attachments := make([]Attachment, n)
11	content := []byte(strings.Repeat("x", contentSize))
12	for i := range n {
13		attachments[i] = Attachment{
14			FilePath: fmt.Sprintf("/path/to/file%d.txt", i),
15			MimeType: "text/plain",
16			Content:  content,
17		}
18	}
19	return attachments
20}
21
22func BenchmarkPromptWithTextAttachments(b *testing.B) {
23	cases := []struct {
24		name        string
25		numFiles    int
26		contentSize int
27	}{
28		{"1file_100bytes", 1, 100},
29		{"5files_1KB", 5, 1024},
30		{"10files_10KB", 10, 10 * 1024},
31		{"20files_50KB", 20, 50 * 1024},
32	}
33
34	for _, tc := range cases {
35		attachments := makeTestAttachments(tc.numFiles, tc.contentSize)
36		prompt := "Process these files"
37
38		b.Run(tc.name, func(b *testing.B) {
39			b.ReportAllocs()
40			for range b.N {
41				_ = PromptWithTextAttachments(prompt, attachments)
42			}
43		})
44	}
45}