git.giacomolli.uk - onechat/blob - terminal_image_test.go

   1 package main
   2
   3 import (
   4         "bytes"
   5         "encoding/json"
   6         "net"
   7         "os"
   8         "path/filepath"
   9         "strings"
  10         "testing"
  11
  12         "github.com/stretchr/testify/require"
  13 )
  14
  15 func TestTerminalImageLineSendsAnAbsolutePathAndCaption(t *testing.T) {
  16         client, daemon := net.Pipe()
  17         defer func() {
  18                 require.NoError(t, client.Close())
  19                 require.NoError(t, daemon.Close())
  20         }()
  21
  22         done := make(chan error, 1)
  23         go func() {
  24                 done <- terminalSendLines(t.Context(), client, strings.NewReader("/image 'image sample.png' Read the label.\n"))
  25         }()
  26
  27         var request SocketRequest
  28
  29         require.NoError(t, json.NewDecoder(daemon).Decode(&request))
  30         require.NoError(t, <-done)
  31
  32         expected, err := filepath.Abs("image sample.png")
  33         require.NoError(t, err)
  34         require.Equal(t, SocketRequest{
  35                 Type: "image",
  36                 Path: expected,
  37                 Text: "Read the label.",
  38         }, request)
  39 }
  40
  41 func TestTerminalImagePathsAndCaptionsHaveNoShellExpansion(t *testing.T) {
  42         testCases := []struct {
  43                 name    string
  44                 line    string
  45                 path    string
  46                 caption string
  47                 command bool
  48         }{
  49                 {
  50                         name:    "accepts paths without captions",
  51                         line:    "/image sample.png",
  52                         path:    "sample.png",
  53                         command: true,
  54                 },
  55                 {
  56                         name:    "keeps spaces in double quoted paths",
  57                         line:    `/image "a b.png" caption`,
  58                         path:    "a b.png",
  59                         caption: "caption",
  60                         command: true,
  61                 },
  62                 {
  63                         name:    "caption stays literal",
  64                         line:    "/image sample.png $(not-a-command)  ",
  65                         path:    "sample.png",
  66                         caption: "$(not-a-command)  ",
  67                         command: true,
  68                 },
  69                 {
  70                         name:    "missing path refused",
  71                         line:    "/image",
  72                         command: true,
  73                 },
  74                 {
  75                         name:    "unclosed quote refused",
  76                         line:    `/image "sample.png`,
  77                         command: true,
  78                 },
  79                 {
  80                         name:    "joined caption refused",
  81                         line:    `/image "sample.png"caption`,
  82                         command: true,
  83                 },
  84                 {
  85                         name:    "other command stays text",
  86                         line:    "/images sample.png",
  87                         command: false,
  88                 },
  89         }
  90
  91         for _, tc := range testCases {
  92                 t.Run(tc.name, func(t *testing.T) {
  93                         request, command := terminalImageRequest(tc.line)
  94                         require.Equal(t, tc.command, command)
  95                         require.Equal(t, tc.path, request.Path)
  96                         require.Equal(t, tc.caption, request.Text)
  97                 })
  98         }
  99 }
 100
 101 func TestTerminalImageIsStoredBeforeItReachesTheModel(t *testing.T) {
 102         server := newModelFake(t)
 103         socket := daemonStartTest(t, server)
 104         terminal := terminalStartTestClient(t, socket)
 105         terminal.WaitFor(t, "[connected to OneChat]")
 106
 107         path := filepath.Join(t.TempDir(), "sample.png")
 108         data := imageTestPNG(t, 2, 2)
 109         require.NoError(t, os.WriteFile(path, data, 0o600))
 110         terminal.Type("/image " + path + " What is this?")
 111         terminal.WaitFor(t, "[done]")
 112         require.Equal(t, 1, server.Calls())
 113         require.Contains(t, string(server.Call(0).raw), "input_image")
 114         require.Contains(t, string(server.Call(0).raw), "What is this?")
 115         require.NotContains(t, string(server.Call(0).raw), path)
 116
 117         files, err := os.ReadDir(filepath.Join(filepath.Dir(socket), imageDirName))
 118         require.NoError(t, err)
 119         require.Len(t, files, 1)
 120
 121         stored, err := os.ReadFile(filepath.Join(filepath.Dir(socket), imageDirName, files[0].Name()))
 122         require.NoError(t, err)
 123         require.Equal(t, data, stored)
 124         require.NotContains(t, terminal.Output(), "base64")
 125 }
 126
 127 func TestTerminalRefusesLargeAndUnsupportedImagesWithoutAModelCall(t *testing.T) {
 128         server := newModelFake(t)
 129         test := newTestTurnLoop(t, server)
 130         path := filepath.Join(t.TempDir(), "sample.png")
 131         require.NoError(t, os.WriteFile(path, bytes.Repeat([]byte{0}, imageMaxBytes+1), 0o600))
 132
 133         request := SocketRequest{
 134                 Type: "image",
 135                 Path: path,
 136         }
 137         require.ErrorContains(t, daemonRunRequest(test.loop, nil, request, func(string) {
 138         }), "exceeds 20 MB")
 139         require.NoError(t, os.WriteFile(path, []byte("not an image"), 0o600))
 140         require.ErrorContains(t, daemonRunRequest(test.loop, nil, request, func(string) {
 141         }), "JPEG, PNG, or WebP")
 142
 143         require.Equal(t, 0, test.memory.Len())
 144         require.Equal(t, 0, server.Calls())
 145
 146         _, err := os.Stat(test.agent.Images().dir)
 147         require.ErrorIs(t, err, os.ErrNotExist)
 148 }
 149
 150 func TestImageCommandErrorDoesNotEndTheTerminalClient(t *testing.T) {
 151         socket := daemonStartTest(t, newModelFake(t))
 152         terminal := terminalStartTestClient(t, socket)
 153         terminal.WaitFor(t, "[connected to OneChat]")
 154         terminal.Type("/image")
 155         terminal.WaitFor(t, "use /image PATH [caption]")
 156         terminal.Type("Continue.")
 157         terminal.WaitFor(t, "[done]")
 158         require.NoError(t, terminal.Close())
 159 }