From 0710906fc7467385f5f23cd7e18e3ed61d3706e7 Mon Sep 17 00:00:00 2001
From: Lei Kang <lkang@nuro.ai>
Date: Fri, 12 Sep 2025 16:32:50 -0700
Subject: [PATCH] fix the test

---
 pkg/codec/vpx/vpx_decoder.go |  3 +-
 pkg/codec/vpx/vpx_test.go    | 66 +++++++++++++++++++++---------------
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/pkg/codec/vpx/vpx_decoder.go b/pkg/codec/vpx/vpx_decoder.go
index a34d1d9..97e8492 100644
--- a/pkg/codec/vpx/vpx_decoder.go
+++ b/pkg/codec/vpx/vpx_decoder.go
@@ -99,7 +99,6 @@ func NewDecoder(r io.Reader, p prop.Media) (codec.VideoDecoder, error) {
 }
 
 func (d *decoder) Read() (image.Image, func(), error) {
-
 	n, err := d.reader.Read(d.buf)
 	if err != nil {
 		return nil, nil, err
@@ -107,7 +106,7 @@ func (d *decoder) Read() (image.Image, func(), error) {
 	if n > 0 {
 		status := C.decodeFrame(d.codec, (*C.uint8_t)(&d.buf[0]), C.uint(n))
 		if status != C.VPX_CODEC_OK {
-			return nil, nil, fmt.Errorf("Decode failed: %v", status)
+			return nil, nil, fmt.Errorf("decode failed: %v", status)
 		}
 	}
 
diff --git a/pkg/codec/vpx/vpx_test.go b/pkg/codec/vpx/vpx_test.go
index 141500b..bb661d6 100644
--- a/pkg/codec/vpx/vpx_test.go
+++ b/pkg/codec/vpx/vpx_test.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"math"
 	"math/rand"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -449,7 +450,7 @@ func TestVP8EncodeDecode(t *testing.T) {
 		}
 		p.LagInFrames = 0 // Disable frame lag buffering for real-time encoding
 		p.RateControlEndUsage = RateControlCBR
-		totalFrames := 100
+		totalFrames := 10
 		var cnt uint32
 		r, err := p.BuildVideoEncoder(
 			video.ReaderFunc(func() (image.Image, func(), error) {
@@ -473,36 +474,47 @@ func TestVP8EncodeDecode(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		data, rel, err := r.Read()
-		if err != nil {
-			t.Fatal(err)
-		}
-		defer rel()
+		var wg sync.WaitGroup
+		wg.Add(1)
 
-		// Decode the frame
-		writer.Write(data)
-		writer.Close()
-
-		// Poll for frame with timeout
-		timeout := time.After(2 * time.Second)
-		ticker := time.NewTicker(10 * time.Millisecond)
-		defer ticker.Stop()
-
-		for {
-			select {
-			case <-timeout:
-				t.Fatal("Timeout: No frame received within 2 seconds")
-			case <-ticker.C:
-				frame, rel, err := decoder.Read()
+		counter := 0
+		go func() {
+			defer wg.Done()
+			for {
+				img, rel, err := decoder.Read()
+				if err == io.EOF {
+					return
+				}
 				if err != nil {
-					t.Fatal(err)
-				}
-				defer rel()
-				if frame != nil {
-					t.Log("Successfully received and decoded frame")
-					return // Test passes
+					t.Errorf("decoder read error: %v", err)
+					return
 				}
+				assert.Equal(t, initialWidth, img.Bounds().Dx())
+				assert.Equal(t, initialHeight, img.Bounds().Dy())
+				rel()
+				counter++
+			}
+		}()
+
+		// --- feed encoded frames to writer
+		for {
+			data, rel, err := r.Read()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				t.Fatalf("encoder error: %v", err)
+			}
+			_, werr := writer.Write(data)
+			rel()
+			if werr != nil {
+				t.Fatalf("writer error: %v", werr)
 			}
 		}
+		writer.Close()
+
+		// ✅ wait until decoder goroutine is done
+		wg.Wait()
+		assert.Equal(t, counter, totalFrames)
 	})
 }