TeleportFling/internal/protocol/jpeg_test.go

package protocol

import (
	"bytes"
	"image"
	"image/color"
	"testing"
)

// minimalRGB builds a tiny RGBA test image with a known gradient.
func minimalRGB(w, h int) *image.RGBA {
	img := image.NewRGBA(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			img.SetRGBA(x, y, color.RGBA{R: uint8(x * 4), G: uint8(y * 4), B: 128, A: 255})
		}
	}
	return img
}

// minimalYCbCr builds a 420 subsampled image with a vertical colour split.
func minimalYCbCr(w, h int) *image.YCbCr {
	img := image.NewYCbCr(image.Rect(0, 0, w, h), image.YCbCrSubsampleRatio420)
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			r, g, b := uint8(0), uint8(0), uint8(0)
			switch {
			case x < w/3:
				r = 200
			case x < w*2/3:
				g = 200
			default:
				b = 200
			}
			img.Y[y*img.YStride+x], _, _ = color.RGBToYCbCr(r, g, b)
		}
	}
	// Subsample chroma from the known dominant colours per region.
	for y := 0; y < h/2; y++ {
		for x := 0; x < w/2; x++ {
			idx := (y*img.CStride + x)
			switch {
			case x < w/6:
				_, cb, cr := color.RGBToYCbCr(200, 0, 0)
				img.Cb[idx], img.Cr[idx] = cb, cr
			case x < w*2/6:
				_, cb, cr := color.RGBToYCbCr(0, 200, 0)
				img.Cb[idx], img.Cr[idx] = cb, cr
			default:
				_, cb, cr := color.RGBToYCbCr(0, 0, 200)
				img.Cb[idx], img.Cr[idx] = cb, cr
			}
		}
	}
	return img
}

// TestJPEGEncodeRGBA checks the RGBA path produces a decodable JPEG.
func TestJPEGEncodeRGBA(t *testing.T) {
	enc, err := NewJPEGEncoder()
	if err != nil {
		t.Fatalf("NewJPEGEncoder: %v", err)
	}
	defer enc.Close()

	img := minimalRGB(64, 48)
	buf, err := enc.Encode(img, 80)
	if err != nil {
		t.Fatalf("Encode: %v", err)
	}
	if len(buf) == 0 {
		t.Fatal("empty JPEG output")
	}
	if !bytes.Equal(buf[:2], []byte{0xFF, 0xD8}) {
		t.Errorf("bad JPEG SOI marker: %x", buf[:2])
	}
	if !bytes.Equal(buf[len(buf)-2:], []byte{0xFF, 0xD9}) {
		t.Errorf("bad JPEG EOI marker")
	}
}

// TestJPEGEncodeYCbCr checks the YUV path for all three subsampling modes.
func TestJPEGEncodeYCbCr(t *testing.T) {
	enc, err := NewJPEGEncoder()
	if err != nil {
		t.Fatalf("NewJPEGEncoder: %v", err)
	}
	defer enc.Close()

	img := minimalYCbCr(64, 48)
	if _, err := enc.Encode(img, 80); err != nil {
		t.Fatalf("Encode (420): %v", err)
	}

	img422 := minimalYCbCr(64, 48)
	img422.SubsampleRatio = image.YCbCrSubsampleRatio422
	if _, err := enc.Encode(img422, 80); err != nil {
		t.Fatalf("Encode (422): %v", err)
	}

	img444 := minimalYCbCr(64, 48)
	img444.SubsampleRatio = image.YCbCrSubsampleRatio444
	if _, err := enc.Encode(img444, 80); err != nil {
		t.Fatalf("Encode (444): %v", err)
	}
}

// TestJPEGEncodeQualityClamp ensures out-of-range qualities are clamped.
func TestJPEGEncodeQualityClamp(t *testing.T) {
	enc, err := NewJPEGEncoder()
	if err != nil {
		t.Fatalf("NewJPEGEncoder: %v", err)
	}
	defer enc.Close()

	img := minimalRGB(16, 16)
	for _, q := range []int{-5, 0, 101, 200} {
		if buf, err := enc.Encode(img, q); err != nil || len(buf) == 0 {
			t.Errorf("quality %d: err=%v len=%d", q, err, len(buf))
		}
	}
}

// TestJPEGEncodeBGRA checks the raw-BGRA fast path produces a decodable
// JPEG with correct SOI/EOI markers.
func TestJPEGEncodeBGRA(t *testing.T) {
	enc, err := NewJPEGEncoder()
	if err != nil {
		t.Fatalf("NewJPEGEncoder: %v", err)
	}
	defer enc.Close()

	// 16x8 BGRA: top half red (R,G,B=255,0,0), bottom half blue (0,0,255).
	w, h := 16, 8
	pix := make([]byte, w*h*4)
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			off := (y*w + x) * 4
			if y < h/2 {
				pix[off], pix[off+1], pix[off+2], pix[off+3] = 0, 0, 255, 255 // B,G,R,A = blue,red
			} else {
				pix[off], pix[off+1], pix[off+2], pix[off+3] = 255, 0, 0, 255 // B,G,R,A = red,blue
			}
		}
	}

	buf, err := enc.EncodeBGRA(pix, w, h, 85)
	if err != nil {
		t.Fatalf("EncodeBGRA: %v", err)
	}
	if len(buf) == 0 {
		t.Fatal("empty JPEG output")
	}
	if !bytes.Equal(buf[:2], []byte{0xFF, 0xD8}) {
		t.Errorf("bad JPEG SOI marker: %x", buf[:2])
	}
	if !bytes.Equal(buf[len(buf)-2:], []byte{0xFF, 0xD9}) {
		t.Errorf("bad JPEG EOI marker")
	}
}

// TestJPEGEncodeBGRAShort rejects truncated pixel buffers.
func TestJPEGEncodeBGRAShort(t *testing.T) {
	enc, err := NewJPEGEncoder()
	if err != nil {
		t.Fatalf("NewJPEGEncoder: %v", err)
	}
	defer enc.Close()

	if _, err := enc.EncodeBGRA(make([]byte, 10), 16, 8, 80); err == nil {
		t.Error("short BGRA buffer accepted")
	}
}