package attach
import (
"bytes"
"encoding/base64"
"errors"
"image"
"image/color"
"image/png"
"os"
"path/filepath"
"strings"
"testing"
"codey/internal/llm"
)
// tinyPNG is a real one-pixel PNG, encoded rather than pasted in as a literal
// so the magic bytes under test are the ones a library actually writes.
func tinyPNG(t *testing.T) []byte {
t.Helper()
im := image.NewRGBA(image.Rect(0, 0, 1, 1))
im.Set(0, 0, color.RGBA{R: 1, G: 2, B: 3, A: 255})
var b bytes.Buffer
if err := png.Encode(&b, im); err != nil {
t.Fatal(err)
}
return b.Bytes()
}
func TestSniffKnowsTheFourFormats(t *testing.T) {
cases := []struct {
name string
data []byte
want string
}{
{"png", tinyPNG(t), "image/png"},
{"jpeg", []byte{0xff, 0xd8, 0xff, 0xe0, 0}, "image/jpeg"},
{"gif", []byte("GIF89a....."), "image/gif"},
{"webp", []byte("RIFF\x00\x00\x00\x00WEBPVP8 "), "image/webp"},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
got, ok := Sniff(c.data)
if !ok || got != c.want {
t.Errorf("Sniff = %q, %v; want %q, true", got, ok, c.want)
}
})
}
}
// Text on the clipboard is the ordinary case for ctrl-V, and it must not come
// back looking like a broken image.
func TestSniffRejectsNonImages(t *testing.T) {
for _, s := range []string{"", "hello", "<svg></svg>", "%PDF-1.4"} {
if got, ok := Sniff([]byte(s)); ok {
t.Errorf("Sniff(%q) = %q, true; want false", s, got)
}
}
}
// An extension is a claim, not evidence: the bytes decide.
func TestNewRejectsAMislabelledFile(t *testing.T) {
_, err := New([]byte("this is just text"), "notes.png")
if err == nil {
t.Fatal("text passed as an image")
}
if errors.Is(err, ErrTooBig) {
t.Errorf("wrong complaint: %v", err)
}
}
func TestNewRefusesAnOversizeImage(t *testing.T) {
big := append(tinyPNG(t), make([]byte, MaxBytes)...)
_, err := New(big, "huge.png")
if !errors.Is(err, ErrTooBig) {
t.Fatalf("err = %v; want ErrTooBig", err)
}
// The message has to say how big, because the next move is to shrink it.
if !strings.Contains(err.Error(), "MB") {
t.Errorf("message does not give a size: %v", err)
}
}
func TestBlockCarriesBase64OfTheBytes(t *testing.T) {
data := tinyPNG(t)
im, err := New(data, "x.png")
if err != nil {
t.Fatal(err)
}
blk := im.Block()
if blk.Type != llm.BlockImage || blk.Source == nil {
t.Fatalf("block = %+v", blk)
}
if blk.Source.Type != "base64" || blk.Source.MediaType != "image/png" {
t.Errorf("source = %+v", blk.Source)
}
got, err := base64.StdEncoding.DecodeString(blk.Source.Data)
if err != nil || !bytes.Equal(got, data) {
t.Errorf("data did not survive encoding: %v", err)
}
}
// ---- paths --------------------------------------------------------------
func TestPathAcceptsWhatATerminalPastes(t *testing.T) {
dir := t.TempDir()
plain := filepath.Join(dir, "shot.png")
spaced := filepath.Join(dir, "my holiday.png")
for _, p := range []string{plain, spaced} {
if err := os.WriteFile(p, tinyPNG(t), 0o644); err != nil {
t.Fatal(err)
}
}
cases := []struct {
name string
paste string
want string
}{
{"bare", plain, plain},
{"trailing newline", plain + "\n", plain},
{"quoted", `"` + spaced + `"`, spaced},
{"shell escaped", strings.ReplaceAll(spaced, " ", `\ `), spaced},
{"file url", "file://" + strings.ReplaceAll(spaced, " ", "%20"), spaced},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
got, ok := Path(c.paste)
if !ok || got != c.want {
t.Errorf("Path(%q) = %q, %v; want %q, true", c.paste, got, ok, c.want)
}
})
}
}
// Text that merely mentions a file is text. Turning it into an attachment
// would silently eat what somebody typed.
func TestPathDeclinesOrdinaryText(t *testing.T) {
dir := t.TempDir()
real := filepath.Join(dir, "shot.png")
if err := os.WriteFile(real, tinyPNG(t), 0o644); err != nil {
t.Fatal(err)
}
for _, s := range []string{
"look at " + real, // a sentence, not a path
real + "\nand the other one", // more than one line
filepath.Join(dir, "missing.png"), // nothing there
filepath.Join(dir, "notes.txt"), // not an image
dir, // a directory
"",
} {
if got, ok := Path(s); ok {
t.Errorf("Path(%q) = %q, true; want false", s, got)
}
}
}
func TestFromFileReadsAnImage(t *testing.T) {
p := filepath.Join(t.TempDir(), "shot.png")
if err := os.WriteFile(p, tinyPNG(t), 0o644); err != nil {
t.Fatal(err)
}
im, err := FromFile(p)
if err != nil {
t.Fatal(err)
}
if im.Media != "image/png" || im.Origin != "shot.png" {
t.Errorf("image = %q from %q", im.Media, im.Origin)
}
}
// ---- markers ------------------------------------------------------------
func fill(t *testing.T, s *Set, n int) {
t.Helper()
for i := 0; i < n; i++ {
im, err := New(tinyPNG(t), "clipboard")
if err != nil {
t.Fatal(err)
}
if got := s.Add(im); got != i+1 {
t.Fatalf("Add returned %d, want %d", got, i+1)
}
}
}
// kinds renders a block list as a shape — "text,image,text" — because what
// matters about the expansion is the order.
func kinds(blocks []llm.ContentBlock) string {
var out []string
for _, b := range blocks {
out = append(out, b.Type)
}
return strings.Join(out, ",")
}
func TestBlocksPutsThePictureWhereTheMarkerStood(t *testing.T) {
var s Set
fill(t, &s, 1)
blocks := s.Blocks("here is the bug [image 1] can you see it")
if got := kinds(blocks); got != "text,image,text" {
t.Fatalf("shape = %s", got)
}
if blocks[0].Text != "here is the bug" {
t.Errorf("leading text = %q", blocks[0].Text)
}
if blocks[2].Text != "can you see it" {
t.Errorf("trailing text = %q", blocks[2].Text)
}
}
func TestBlocksKeepsTwoImagesInOrder(t *testing.T) {
var s Set
fill(t, &s, 2)
// Deliberately out of order in the text: the marker says which is which,
// not the position.
blocks := s.Blocks("after [image 2] before [image 1]")
if got := kinds(blocks); got != "text,image,text,image" {
t.Fatalf("shape = %s", got)
}
if blocks[1].Source.Data == "" || blocks[3].Source.Data == "" {
t.Error("an image block lost its data")
}
}
// Deleting the marker is the undo. The image it named must not be sent
// anyway, tacked on somewhere the model cannot place it.
func TestBlocksDropsAnImageWhoseMarkerWasDeleted(t *testing.T) {
var s Set
fill(t, &s, 2)
blocks := s.Blocks("only this one please [image 2]")
if got := kinds(blocks); got != "text,image" {
t.Fatalf("shape = %s", got)
}
}
// A number nobody pasted is just something typed, and typing it should reach
// the model as written.
func TestBlocksLeavesAnUnknownMarkerAsText(t *testing.T) {
var s Set
fill(t, &s, 1)
blocks := s.Blocks("what does [image 7] mean")
if got := kinds(blocks); got != "text" {
t.Fatalf("shape = %s", got)
}
if blocks[0].Text != "what does [image 7] mean" {
t.Errorf("text = %q", blocks[0].Text)
}
}
func TestBlocksGivesAWordlessMessageWords(t *testing.T) {
var s Set
fill(t, &s, 1)
blocks := s.Blocks("[image 1]")
if got := kinds(blocks); got != "image,text" {
t.Fatalf("shape = %s", got)
}
if strings.TrimSpace(blocks[1].Text) == "" {
t.Error("the message has no text at all")
}
}
func TestBlocksOfPlainTextIsOneTextBlock(t *testing.T) {
var s Set
blocks := s.Blocks("just a question")
if got := kinds(blocks); got != "text" {
t.Fatalf("shape = %s", got)
}
}
func TestResetEmptiesTheSet(t *testing.T) {
var s Set
fill(t, &s, 2)
s.Reset()
if s.Len() != 0 {
t.Fatalf("Len = %d after Reset", s.Len())
}
// Numbering starts over with the next prompt.
fill(t, &s, 1)
}
func TestMarkerReads(t *testing.T) {
if got := Marker(3); got != "[image 3]" {
t.Errorf("Marker(3) = %q", got)
}
}
// ---- paths left in the prompt -------------------------------------------
// The case this exists for: a terminal that does not bracket its pastes drops
// the path in as if it had been typed, and nothing catches it until the
// prompt is sent.
func TestExpandPathsTakesAPastedPath(t *testing.T) {
p := filepath.Join(t.TempDir(), "clip.png")
if err := os.WriteFile(p, tinyPNG(t), 0o644); err != nil {
t.Fatal(err)
}
var s Set
got := s.ExpandPaths("what is wrong here " + p)
if got != "what is wrong here [image 1]" {
t.Fatalf("text = %q", got)
}
if s.Len() != 1 {
t.Fatalf("%d images attached", s.Len())
}
}
// Terminals hand over a dropped or copied file quoted. The quotes go with the
// path, not into the sentence.
func TestExpandPathsTakesAQuotedPathAlone(t *testing.T) {
p := filepath.Join(t.TempDir(), "clip.png")
if err := os.WriteFile(p, tinyPNG(t), 0o644); err != nil {
t.Fatal(err)
}
var s Set
if got := s.ExpandPaths(`"` + p + `"`); got != "[image 1]" {
t.Fatalf("text = %q", got)
}
}
func TestExpandPathsTakesTwo(t *testing.T) {
dir := t.TempDir()
var paths []string
for _, n := range []string{"a.png", "b.png"} {
p := filepath.Join(dir, n)
if err := os.WriteFile(p, tinyPNG(t), 0o644); err != nil {
t.Fatal(err)
}
paths = append(paths, p)
}
var s Set
got := s.ExpandPaths("compare " + paths[0] + " with " + paths[1])
if got != "compare [image 1] with [image 2]" {
t.Fatalf("text = %q", got)
}
if s.Len() != 2 {
t.Fatalf("%d images attached", s.Len())
}
}
// A prompt about the code must come through untouched. This is the risk the
// feature carries, so it is the test that matters most.
func TestExpandPathsLeavesOrdinaryPromptsAlone(t *testing.T) {
dir := t.TempDir()
missing := filepath.Join(dir, "gone.png")
source := filepath.Join(dir, "model.go")
if err := os.WriteFile(source, []byte("package tui"), 0o644); err != nil {
t.Fatal(err)
}
for _, prompt := range []string{
"rename shot.png to screenshot.png",
"why does " + source + " not compile",
"the file " + missing + " is missing",
"look at internal/tui/model.go:44",
"no paths at all here",
} {
var s Set
if got := s.ExpandPaths(prompt); got != prompt {
t.Errorf("ExpandPaths(%q) = %q", prompt, got)
}
if s.Len() != 0 {
t.Errorf("%q attached %d images", prompt, s.Len())
}
}
}
// A marker already in the text is left as it is: the paste key put it there
// and its picture is already attached.
func TestExpandPathsLeavesExistingMarkers(t *testing.T) {
var s Set
fill(t, &s, 1)
if got := s.ExpandPaths("look at [image 1] closely"); got != "look at [image 1] closely" {
t.Fatalf("text = %q", got)
}
if s.Len() != 1 {
t.Fatalf("%d images after expanding", s.Len())
}
}
// The same file named twice is one picture, and both mentions point at it.
func TestExpandPathsMentionedTwiceAttachesOnce(t *testing.T) {
p := filepath.Join(t.TempDir(), "clip.png")
if err := os.WriteFile(p, tinyPNG(t), 0o644); err != nil {
t.Fatal(err)
}
var s Set
got := s.ExpandPaths(p + " and again " + p)
if s.Len() != 1 {
t.Fatalf("%d images attached, want 1", s.Len())
}
if got != "[image 1] and again [image 1]" {
t.Fatalf("text = %q", got)
}
}