diff --git a/container/mts/encoder_test.go b/container/mts/encoder_test.go
index 47e9a809..6b46703e 100644
--- a/container/mts/encoder_test.go
+++ b/container/mts/encoder_test.go
@@ -101,7 +101,10 @@ func TestEncodeVideo(t *testing.T) {
 
 	// Create the dst and write the test data to encoder.
 	dst := &destination{}
-	_, err := NewEncoder(nopCloser{dst}, 25, EncodeH264).Write(data)
+	e := NewEncoder(nopCloser{dst}, 25, EncodeH264)
+	e.NALBasedPSI(false, psiSendCount)
+
+	_, err := e.Write(data)
 	if err != nil {
 		t.Fatalf("could not write data to encoder, failed with err: %v\n", err)
 	}