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         for _, tc := range testCases {
  91                 t.Run(tc.name, func(t *testing.T) {
  92                         request, command := terminalImageRequest(tc.line)
  93                         require.Equal(t, tc.command, command)
  94                         require.Equal(t, tc.path, request.Path)
  95                         require.Equal(t, tc.caption, request.Text)
  96                 })
  97         }
  98 }
  99
 100 func TestTerminalImageIsStoredBeforeItReachesTheModel(t *testing.T) {
 101         server := newModelFake(t)
 102         socket := daemonStartTest(t, server)
 103         terminal := terminalStartTestClient(t, socket)
 104         terminal.WaitFor(t, "[connected to OneChat]")
 105
 106         path := filepath.Join(t.TempDir(), "sample.png")
 107         data := imageTestPNG(t, 2, 2)
 108         require.NoError(t, os.WriteFile(path, data, 0o600))
 109         terminal.Type("/image " + path + " What is this?")
 110         terminal.WaitFor(t, "[done]")
 111         require.Equal(t, 1, server.Calls())
 112         require.Contains(t, string(server.Call(0).raw), "input_image")
 113         require.Contains(t, string(server.Call(0).raw), "What is this?")
 114         require.NotContains(t, string(server.Call(0).raw), path)
 115
 116         files, err := os.ReadDir(filepath.Join(filepath.Dir(socket), imageDirName))
 117         require.NoError(t, err)
 118         require.Len(t, files, 1)
 119
 120         stored, err := os.ReadFile(filepath.Join(filepath.Dir(socket), imageDirName, files[0].Name()))
 121         require.NoError(t, err)
 122         require.Equal(t, data, stored)
 123         require.NotContains(t, terminal.Output(), "base64")
 124 }
 125
 126 func TestTerminalRefusesLargeAndUnsupportedImagesWithoutAModelCall(t *testing.T) {
 127         server := newModelFake(t)
 128         test := newTestTurnLoop(t, server)
 129         path := filepath.Join(t.TempDir(), "sample.png")
 130         require.NoError(t, os.WriteFile(path, bytes.Repeat([]byte{0}, imageMaxBytes+1), 0o600))
 131
 132         request := SocketRequest{
 133                 Type: "image",
 134                 Path: path,
 135         }
 136         require.ErrorContains(t, daemonRunRequest(test.loop, nil, request, func(string) {
 137         }), "exceeds 20 MB")
 138         require.NoError(t, os.WriteFile(path, []byte("not an image"), 0o600))
 139         require.ErrorContains(t, daemonRunRequest(test.loop, nil, request, func(string) {
 140         }), "JPEG, PNG, or WebP")
 141         require.Equal(t, 0, test.memory.Len())
 142         require.Equal(t, 0, server.Calls())
 143
 144         _, err := os.Stat(test.agent.Images().dir)
 145         require.ErrorIs(t, err, os.ErrNotExist)
 146 }
 147
 148 func TestImageCommandErrorDoesNotEndTheTerminalClient(t *testing.T) {
 149         socket := daemonStartTest(t, newModelFake(t))
 150         terminal := terminalStartTestClient(t, socket)
 151         terminal.WaitFor(t, "[connected to OneChat]")
 152         terminal.Type("/image")
 153         terminal.WaitFor(t, "use /image PATH [caption]")
 154         terminal.Type("Continue.")
 155         terminal.WaitFor(t, "[done]")
 156         require.NoError(t, terminal.Close())
 157 }