1package hints
2
3import (
4 "bytes"
5 "context"
6 "encoding/binary"
7 "fmt"
8 "io"
9 "math"
10 "net"
11 "os"
12 "path/filepath"
13 "testing"
14 "testing/synctest"
15 "time"
16
17 "github.com/stretchr/testify/require"
18)
19
20// u32, u64 and f32 build the byte fixtures of the protocol.
21func u32(v uint32) []byte { return binary.LittleEndian.AppendUint32(nil, v) }
22func u64(v uint64) []byte { return binary.LittleEndian.AppendUint64(nil, v) }
23func f32(v float32) []byte { return u32(math.Float32bits(v)) }
24func cat(parts ...[]byte) []byte {
25 return bytes.Join(parts, nil)
26}
27
28func TestWriteRequest(t *testing.T) {
29 var buf bytes.Buffer
30 pix := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}
31 require.NoError(t, writeRequest(&buf, 0x0102030405, detFrame{width: 2, height: 1, stride: 12, m: &mapping{b: append(pix, 99)}}))
32 want := cat([]byte("NHRQ"), u32(1), u64(0x0102030405), u32(2), u32(1), u32(12), u32(1), u32(0), f32(0), pix)
33 require.Equal(t, want, buf.Bytes())
34}
35
36func TestWriteRequestRegions(t *testing.T) {
37 var buf bytes.Buffer
38 pix := []byte{1, 2, 3, 4}
39 f := detFrame{width: 1, height: 1, stride: 4, m: &mapping{b: pix}, regions: []region{{0, 1, 2, 3}, {4, 5, 6, 7}}}
40 require.NoError(t, writeRequest(&buf, 1, f))
41 want := cat([]byte("NHRQ"), u32(1), u64(1), u32(1), u32(1), u32(4), u32(1), u32(2),
42 u32(0), u32(1), u32(2), u32(3), u32(4), u32(5), u32(6), u32(7), f32(0), pix)
43 require.Equal(t, want, buf.Bytes())
44}
45
46func TestReadHello(t *testing.T) {
47 dev, err := readHello(bytes.NewReader(cat([]byte("NHNT"), u32(1), u32(0), u32(6), []byte("Vulkan"), u32(1280))))
48 require.NoError(t, err)
49 require.Equal(t, "Vulkan", dev)
50
51 _, err = readHello(bytes.NewReader(cat([]byte("NHNT"), u32(1), u32(1), u32(4), []byte("boom"))))
52 require.ErrorIs(t, err, errDetector)
53 require.ErrorContains(t, err, "boom")
54
55 _, err = readHello(bytes.NewReader(cat([]byte("XXXX"), u32(1))))
56 require.ErrorIs(t, err, errProtocol)
57 _, err = readHello(bytes.NewReader(cat([]byte("NHNT"), u32(2), u32(0))))
58 require.ErrorIs(t, err, errProtocol)
59 _, err = readHello(bytes.NewReader(cat([]byte("NHNT"), u32(1), u32(0), u32(1<<20))))
60 require.ErrorIs(t, err, errProtocol)
61 _, err = readHello(bytes.NewReader([]byte("NHNT")))
62 require.Error(t, err)
63}
64
65func TestReadResponse(t *testing.T) {
66 head := func(seq uint64, status uint32) []byte { return cat([]byte("NHRS"), u32(1), u64(seq), u32(status)) }
67
68 ok := cat(head(5, 0), f32(12.5), u32(2),
69 f32(1), f32(2), f32(30), f32(40), f32(0.9), u32(0),
70 f32(5), f32(6), f32(7), f32(8), f32(0.25), u32(3))
71 got, err := readResponse(bytes.NewReader(ok), 5)
72 require.NoError(t, err)
73 require.Equal(t, []detection{{x1: 1, y1: 2, x2: 30, y2: 40, score: 0.9}, {x1: 5, y1: 6, x2: 7, y2: 8, score: 0.25, region: 3}}, got)
74
75 got, err = readResponse(bytes.NewReader(cat(head(5, 0), f32(1), u32(0))), 5)
76 require.NoError(t, err)
77 require.Empty(t, got)
78
79 _, err = readResponse(bytes.NewReader(cat(head(5, 1), u32(3), []byte("bad"))), 5)
80 require.ErrorIs(t, err, errDetector)
81 require.ErrorContains(t, err, "bad")
82
83 _, err = readResponse(bytes.NewReader(ok), 6) // wrong sequence
84 require.ErrorIs(t, err, errProtocol)
85 _, err = readResponse(bytes.NewReader(cat([]byte("NHRX"), u32(1))), 5)
86 require.ErrorIs(t, err, errProtocol)
87 _, err = readResponse(bytes.NewReader(cat([]byte("NHRS"), u32(9), u64(5), u32(0))), 5)
88 require.ErrorIs(t, err, errProtocol)
89 _, err = readResponse(bytes.NewReader(cat(head(5, 0), f32(1), u32(1<<30))), 5)
90 require.ErrorIs(t, err, errProtocol)
91 _, err = readResponse(bytes.NewReader(ok[:len(ok)-3]), 5) // truncated
92 require.Error(t, err)
93}
94
95func TestLineLog(t *testing.T) {
96 var lines []string
97 l := &lineLog{log: func(s string) { lines = append(lines, s) }}
98 _, _ = l.Write([]byte("one\ntw"))
99 _, _ = l.Write([]byte("o\n\nthree"))
100 require.Equal(t, []string{"one", "two"}, lines)
101}
102
103// The real detector, when the environment names it:
104// NEFERWL_TEST_DETECTOR (binary) and NEFERWL_TEST_DETECTOR_MODEL (models dir).
105func TestProcDetectorIntegration(t *testing.T) {
106 bin, model := os.Getenv("NEFERWL_TEST_DETECTOR"), os.Getenv("NEFERWL_TEST_DETECTOR_MODEL")
107 if bin == "" || model == "" {
108 t.Skip("NEFERWL_TEST_DETECTOR and NEFERWL_TEST_DETECTOR_MODEL not set")
109 }
110 d := newProcDetector(bin, model, 0.05)
111 t.Cleanup(func() { _ = d.Close() })
112 m := &mapping{b: make([]byte, 1280*720*4), live: new(int)}
113 for i := range m.b {
114 m.b[i] = 0xf0 // a blank light frame
115 }
116 f := detFrame{width: 1280, height: 720, stride: 1280 * 4, m: m}
117 ctx := context.Background()
118
119 dets, err := d.Detect(ctx, f)
120 require.NoError(t, err)
121 require.Empty(t, dets) // nothing clickable on a blank frame
122 first := d.proc
123 require.NotNil(t, first)
124 _, err = d.Detect(ctx, f)
125 require.NoError(t, err)
126 require.Same(t, first, d.proc, "the child is kept alive between scans")
127
128 // A dead child fails one scan, and the next one starts a new child.
129 require.NoError(t, first.cmd.Process.Kill())
130 time.Sleep(50 * time.Millisecond)
131 _, err = d.Detect(ctx, f)
132 require.Error(t, err)
133 require.Nil(t, d.proc)
134 _, err = d.Detect(ctx, f)
135 require.NoError(t, err)
136 require.NotSame(t, first, d.proc)
137
138 // A wrong geometry is the detector's own error: the child stays.
139 proc := d.proc
140 _, err = d.Detect(ctx, detFrame{width: 1280, height: 720, stride: 4, m: m})
141 require.ErrorIs(t, err, errDetector)
142 require.Same(t, proc, d.proc)
143 require.NoError(t, d.Close())
144}
145
146func TestProcDetectorMissingBinary(t *testing.T) {
147 d := newProcDetector("/nonexistent/nefer-hints", t.TempDir(), 0.05)
148 _, err := d.Detect(context.Background(), detFrame{width: 1, height: 1, stride: 4, m: &mapping{b: make([]byte, 4)}})
149 require.Error(t, err)
150 require.Nil(t, d.proc)
151 require.NoError(t, d.Close())
152}
153
154func TestProcDetectorCommand(t *testing.T) {
155 for _, threshold := range []float64{0, 0.3} {
156 d := newProcDetector("nefer-hints", "/models", threshold)
157 want := []string{"nefer-hints", "--model", "/models", "--imgsz", "1280", "--gpu", "auto"}
158 if threshold != 0 {
159 want = append(want, "--conf", "0.3")
160 }
161 require.Equal(t, want, d.cmd().Args)
162 }
163}
164
165func TestResponseBoxLimit(t *testing.T) {
166 _, err := readResponse(bytes.NewReader(cat([]byte("NHRS"), u32(1), u64(1), u32(0), f32(0), u32(64001))), 1)
167 require.ErrorIs(t, err, errProtocol)
168}
169
170func TestPrepareDetections(t *testing.T) {
171 regions := []region{{0, 0, 100, 100}, {50, 0, 100, 100}}
172 valid := detection{x1: 90, y1: 10, x2: 100, y2: 30, score: .5}
173 invalid := detection{x1: 80, y1: 0, x2: 120, y2: 40, score: float32(math.NaN()), region: 1}
174 badRegion := valid
175 badRegion.region = 2
176 got := prepareDetections([]detection{valid, invalid, badRegion}, regions, 150, 100)
177 require.Equal(t, []detection{valid}, dropCutBoxes(got, regions))
178 many := make([]detection, maxWorkBoxes+5)
179 for i := range many {
180 many[i] = valid
181 many[i].score = float32(i) / float32(len(many))
182 }
183 got = prepareDetections(many, nil, 150, 100)
184 require.Len(t, got, maxWorkBoxes)
185 require.Equal(t, many[len(many)-1].score, got[0].score)
186 require.Equal(t, many[5].score, got[len(got)-1].score)
187}
188
189// TestDetectorChild is a real subprocess speaking just enough protocol to
190// synchronize cancellation without timers or project-interface doubles.
191func TestDetectorChild(t *testing.T) {
192 mode := os.Getenv("HINT_CHILD_MODE")
193 if mode == "" {
194 return
195 }
196 conn, err := net.Dial("unix", os.Getenv("HINT_CHILD_SOCKET"))
197 if err != nil {
198 os.Exit(2)
199 }
200 defer conn.Close()
201 if mode != "hello" {
202 _, _ = os.Stdout.Write(cat([]byte("NHNT"), u32(1), u32(0), u32(3), []byte("cpu"), u32(1280)))
203 }
204 if mode == "request" {
205 b := make([]byte, 44) // request header, confidence and one BGRA pixel
206 if _, err := io.ReadFull(os.Stdin, b); err != nil {
207 os.Exit(3)
208 }
209 }
210 _, _ = conn.Write([]byte{1})
211 _, _ = io.Copy(io.Discard, os.Stdin)
212 os.Exit(0)
213}
214
215func childDetector(t *testing.T, mode string) (*procDetector, net.Listener) {
216 t.Helper()
217 dir := t.TempDir()
218 socket := filepath.Join(dir, "ready")
219 listener, err := net.Listen("unix", socket)
220 require.NoError(t, err)
221 t.Cleanup(func() { _ = listener.Close() })
222 t.Setenv("HINT_CHILD_MODE", mode)
223 t.Setenv("HINT_CHILD_SOCKET", socket)
224 exe, err := os.Executable()
225 require.NoError(t, err)
226 script := filepath.Join(dir, "detector")
227 require.NoError(t, os.WriteFile(script, fmt.Appendf(nil, "#!/bin/sh\nexec %q -test.run=^TestDetectorChild$\n", exe), 0700))
228 d := newProcDetector(script, dir, 0)
229 t.Cleanup(func() { _ = d.Close() })
230 return d, listener
231}
232
233func TestDetectorCancellation(t *testing.T) {
234 for _, mode := range []string{"hello", "request"} {
235 t.Run(mode, func(t *testing.T) {
236 d, l := childDetector(t, mode)
237 ctx, cancel := context.WithCancel(context.Background())
238 defer cancel()
239 done := make(chan error, 1)
240 go func() {
241 _, err := d.Detect(ctx, detFrame{width: 1, height: 1, stride: 4, m: &mapping{b: make([]byte, 4)}})
242 done <- err
243 }()
244 require.NoError(t, l.(*net.UnixListener).SetDeadline(time.Now().Add(5*time.Second)))
245 conn, err := l.Accept()
246 require.NoError(t, err)
247 require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second)))
248 var ready [1]byte
249 _, err = io.ReadFull(conn, ready[:])
250 _ = conn.Close()
251 require.NoError(t, err)
252 cancel()
253 select {
254 case err := <-done:
255 require.ErrorIs(t, err, context.Canceled)
256 case <-time.After(5 * time.Second):
257 t.Fatal("cancellation did not finish")
258 }
259 })
260 }
261}
262
263func TestDetectorCloseClosesOutput(t *testing.T) {
264 d, l := childDetector(t, "close")
265 require.NoError(t, d.start(context.Background()))
266 conn, err := l.Accept()
267 require.NoError(t, err)
268 _ = conn.Close()
269 p := d.proc
270 require.NoError(t, d.Close())
271 _, err = p.out.Read(make([]byte, 1))
272 require.ErrorIs(t, err, os.ErrClosed)
273 require.NoError(t, d.Close())
274}
275
276// A complete buffered response must not override an already launched context
277// callback, including a detector-reported error that normally keeps the child.
278func TestDetectorCanceledBufferedResponse(t *testing.T) {
279 for _, status := range []uint32{0, 1} {
280 t.Run(fmt.Sprint(status), func(t *testing.T) {
281 d, l := childDetector(t, "close")
282 require.NoError(t, d.start(context.Background()))
283 conn, err := l.Accept()
284 require.NoError(t, err)
285 defer conn.Close()
286 p := d.proc
287 response := cat([]byte("NHRS"), u32(1), u64(1), u32(status))
288 if status == 0 {
289 response = cat(response, f32(0), u32(0))
290 } else {
291 response = cat(response, u32(3), []byte("bad"))
292 }
293 p.rd.Reset(bytes.NewReader(response))
294 ctx, cancel := context.WithCancel(context.Background())
295 cancel()
296 _, err = d.Detect(ctx, detFrame{width: 1, height: 1, stride: 4, m: &mapping{b: make([]byte, 4)}})
297 require.ErrorIs(t, err, context.Canceled)
298 require.Nil(t, d.proc)
299 require.True(t, p.reaped)
300 })
301 }
302}
303
304func TestInterruptDeadlineBoundary(t *testing.T) {
305 for _, cancelContext := range []bool{false, true} {
306 t.Run(fmt.Sprint(cancelContext), func(t *testing.T) {
307 synctest.Test(t, func(t *testing.T) {
308 inR, inW, err := os.Pipe()
309 require.NoError(t, err)
310 defer inR.Close()
311 outR, outW, err := os.Pipe()
312 require.NoError(t, err)
313 defer outW.Close()
314 // Already reaped: only the real pipe endpoints remain to interrupt.
315 p := &process{in: inW, out: outR, reaped: true}
316 ctx, cancel := context.WithCancel(context.Background())
317 defer cancel()
318 stop := interruptOnDeadline(ctx, time.Second, p)
319 if cancelContext {
320 cancel()
321 } else {
322 time.Sleep(time.Second)
323 }
324 synctest.Wait()
325 timedOut, canceled := stop()
326 require.Equal(t, !cancelContext, timedOut)
327 require.Equal(t, cancelContext, canceled)
328 _, err = inW.Write([]byte{1})
329 require.ErrorIs(t, err, os.ErrClosed)
330 _, err = outR.Read(make([]byte, 1))
331 require.ErrorIs(t, err, os.ErrClosed)
332 })
333 })
334 }
335}