wehub-resource-sync
221 lines
7.2 KiB
Go
221 lines
7.2 KiB
Go
package main
|
|
|
|
// main_test.go - end-to-end smoke test for the rfdetr-cpp gRPC backend.
|
|
//
|
|
// Spawns the compiled rfdetr-cpp binary on a free local port, dials it via
|
|
// gRPC, and exercises LoadModel + Detect against the test fixtures
|
|
// downloaded by test.sh. Two scenarios:
|
|
//
|
|
// 1. detection — loads rfdetr-nano-q8_0.gguf and asserts at least one
|
|
// detection comes back with a non-empty class name and a bounding box
|
|
// of non-zero size.
|
|
// 2. segmentation — loads rfdetr-seg-nano-q8_0.gguf and additionally
|
|
// asserts that at least one detection carries a PNG-encoded mask blob
|
|
// (verified by PNG magic bytes).
|
|
//
|
|
// Both specs Skip cleanly if their fixtures are missing so the test target
|
|
// stays usable on a fresh checkout where models haven't been downloaded.
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"fmt"
|
|
"net"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
)
|
|
|
|
func TestRFDetrCpp(t *testing.T) {
|
|
RegisterFailHandler(Fail)
|
|
RunSpecs(t, "rfdetr-cpp backend smoke suite")
|
|
}
|
|
|
|
// freePort grabs an ephemeral TCP port and immediately releases it so the
|
|
// spawned backend can bind to it. There is a tiny TOCTOU window here but in
|
|
// practice it's adequate for a smoke test on a quiet runner.
|
|
func freePort() int {
|
|
l, err := net.Listen("tcp", "127.0.0.1:0")
|
|
Expect(err).ToNot(HaveOccurred(), "freePort listen")
|
|
port := l.Addr().(*net.TCPAddr).Port
|
|
Expect(l.Close()).To(Succeed())
|
|
return port
|
|
}
|
|
|
|
// startBackend spawns the rfdetr-cpp binary on the given port and waits
|
|
// until it accepts TCP connections (up to 10s). The returned cleanup func
|
|
// kills the process and reaps it.
|
|
func startBackend(port int) func() {
|
|
binary, err := filepath.Abs("./rfdetr-cpp")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
if _, err := os.Stat(binary); err != nil {
|
|
Skip(fmt.Sprintf("backend binary not built: %s (run `make rfdetr-cpp` first)", binary))
|
|
}
|
|
|
|
libPath, err := filepath.Abs("./librfdetrcpp-fallback.so")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
if _, err := os.Stat(libPath); err != nil {
|
|
Skip(fmt.Sprintf("fallback library not built: %s (run `make librfdetrcpp-fallback.so` first)", libPath))
|
|
}
|
|
|
|
addr := fmt.Sprintf("127.0.0.1:%d", port)
|
|
cmd := exec.Command(binary, "--addr", addr)
|
|
cmd.Env = append(os.Environ(), "RFDETR_LIBRARY="+libPath)
|
|
cmd.Stdout = os.Stderr
|
|
cmd.Stderr = os.Stderr
|
|
Expect(cmd.Start()).To(Succeed())
|
|
|
|
cleanup := func() {
|
|
if cmd.Process != nil {
|
|
_ = cmd.Process.Kill()
|
|
_, _ = cmd.Process.Wait()
|
|
}
|
|
}
|
|
|
|
deadline := time.Now().Add(10 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
c, err := net.DialTimeout("tcp", addr, 200*time.Millisecond)
|
|
if err == nil {
|
|
_ = c.Close()
|
|
return cleanup
|
|
}
|
|
time.Sleep(200 * time.Millisecond)
|
|
}
|
|
|
|
cleanup()
|
|
Fail(fmt.Sprintf("backend did not become ready on %s within 10s", addr))
|
|
return func() {}
|
|
}
|
|
|
|
// loadTestImage reads the COCO test image downloaded by test.sh and returns
|
|
// its base64-encoded content (the wire format accepted by the Detect RPC).
|
|
func loadTestImage() string {
|
|
imgPath, err := filepath.Abs("test-data/test.jpg")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
imgBytes, err := os.ReadFile(imgPath)
|
|
if err != nil {
|
|
Skip(fmt.Sprintf("test image not present: %s (run test.sh first)", imgPath))
|
|
}
|
|
return base64.StdEncoding.EncodeToString(imgBytes)
|
|
}
|
|
|
|
// dialBackend opens a gRPC client connection to the spawned backend.
|
|
func dialBackend(port int) (pb.BackendClient, func()) {
|
|
addr := fmt.Sprintf("127.0.0.1:%d", port)
|
|
conn, err := grpc.NewClient(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
Expect(err).ToNot(HaveOccurred())
|
|
return pb.NewBackendClient(conn), func() { _ = conn.Close() }
|
|
}
|
|
|
|
// modelPathOrSkip resolves a model file under ./test-models/ and Skip()s
|
|
// the current spec if it's missing.
|
|
func modelPathOrSkip(name string) string {
|
|
modelDir, err := filepath.Abs("test-models")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
modelPath := filepath.Join(modelDir, name)
|
|
if _, err := os.Stat(modelPath); err != nil {
|
|
Skip(fmt.Sprintf("model not present: %s (run test.sh first)", modelPath))
|
|
}
|
|
return modelPath
|
|
}
|
|
|
|
var _ = Describe("rfdetr-cpp backend", func() {
|
|
It("runs object detection against a known-good COCO image", func() {
|
|
modelPath := modelPathOrSkip("rfdetr-nano-q8_0.gguf")
|
|
imgB64 := loadTestImage()
|
|
|
|
port := freePort()
|
|
cleanup := startBackend(port)
|
|
defer cleanup()
|
|
|
|
client, closeConn := dialBackend(port)
|
|
defer closeConn()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
|
defer cancel()
|
|
|
|
loadResp, err := client.LoadModel(ctx, &pb.ModelOptions{
|
|
Model: "rfdetr-nano-q8_0.gguf",
|
|
ModelFile: modelPath,
|
|
Threads: 2,
|
|
})
|
|
Expect(err).ToNot(HaveOccurred(), "LoadModel")
|
|
Expect(loadResp.GetSuccess()).To(BeTrue(), "LoadModel reported failure: %s", loadResp.GetMessage())
|
|
|
|
detResp, err := client.Detect(ctx, &pb.DetectOptions{
|
|
Src: imgB64,
|
|
Threshold: 0.5,
|
|
})
|
|
Expect(err).ToNot(HaveOccurred(), "Detect")
|
|
Expect(detResp.GetDetections()).ToNot(BeEmpty(), "no detections returned on a known-good COCO image")
|
|
|
|
_, _ = fmt.Fprintf(GinkgoWriter, "detection OK: %d detections\n", len(detResp.GetDetections()))
|
|
for i, d := range detResp.GetDetections() {
|
|
Expect(d.GetClassName()).ToNot(BeEmpty(), "detection %d has empty class_name", i)
|
|
Expect(d.GetConfidence()).To(BeNumerically(">=", float32(0.5)),
|
|
"detection %d below threshold", i)
|
|
Expect(d.GetWidth()).To(BeNumerically(">", float32(0)),
|
|
"detection %d has non-positive width", i)
|
|
Expect(d.GetHeight()).To(BeNumerically(">", float32(0)),
|
|
"detection %d has non-positive height", i)
|
|
}
|
|
})
|
|
|
|
It("runs segmentation and returns PNG-encoded masks", func() {
|
|
modelPath := modelPathOrSkip("rfdetr-seg-nano-q8_0.gguf")
|
|
imgB64 := loadTestImage()
|
|
|
|
port := freePort()
|
|
cleanup := startBackend(port)
|
|
defer cleanup()
|
|
|
|
client, closeConn := dialBackend(port)
|
|
defer closeConn()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
|
defer cancel()
|
|
|
|
loadResp, err := client.LoadModel(ctx, &pb.ModelOptions{
|
|
Model: "rfdetr-seg-nano-q8_0.gguf",
|
|
ModelFile: modelPath,
|
|
Threads: 2,
|
|
})
|
|
Expect(err).ToNot(HaveOccurred(), "LoadModel")
|
|
Expect(loadResp.GetSuccess()).To(BeTrue(), "LoadModel reported failure: %s", loadResp.GetMessage())
|
|
|
|
detResp, err := client.Detect(ctx, &pb.DetectOptions{
|
|
Src: imgB64,
|
|
Threshold: 0.5,
|
|
})
|
|
Expect(err).ToNot(HaveOccurred(), "Detect")
|
|
Expect(detResp.GetDetections()).ToNot(BeEmpty(), "no detections returned from segmentation model")
|
|
|
|
haveMask := false
|
|
for i, d := range detResp.GetDetections() {
|
|
m := d.GetMask()
|
|
if len(m) == 0 {
|
|
continue
|
|
}
|
|
haveMask = true
|
|
// Verify PNG magic: 89 50 4E 47 ("\x89PNG").
|
|
Expect(len(m)).To(BeNumerically(">=", 4), "detection %d mask too short", i)
|
|
Expect([]byte{m[0], m[1], m[2], m[3]}).To(Equal([]byte{0x89, 'P', 'N', 'G'}),
|
|
"detection %d mask is not a PNG", i)
|
|
}
|
|
Expect(haveMask).To(BeTrue(),
|
|
"segmentation model returned %d detections but none carried a mask",
|
|
len(detResp.GetDetections()))
|
|
|
|
_, _ = fmt.Fprintf(GinkgoWriter, "segmentation OK: %d detections, at least one with PNG mask\n",
|
|
len(detResp.GetDetections()))
|
|
})
|
|
})
|