github.com/pion/webrtc/v4@v4.0.1/pkg/media/h264writer/h264writer_test.go (about)

     1  // SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
     2  // SPDX-License-Identifier: MIT
     3  
     4  package h264writer
     5  
     6  import (
     7  	"bytes"
     8  	"errors"
     9  	"testing"
    10  
    11  	"github.com/pion/rtp"
    12  	"github.com/stretchr/testify/assert"
    13  )
    14  
    15  type writerCloser struct {
    16  	bytes.Buffer
    17  }
    18  
    19  var errClose = errors.New("close error")
    20  
    21  func (w *writerCloser) Close() error {
    22  	return errClose
    23  }
    24  
    25  func TestNewWith(t *testing.T) {
    26  	writer := &writerCloser{}
    27  	h264Writer := NewWith(writer)
    28  	assert.NotNil(t, h264Writer.Close())
    29  }
    30  
    31  func TestIsKeyFrame(t *testing.T) {
    32  	tests := []struct {
    33  		name    string
    34  		payload []byte
    35  		want    bool
    36  	}{
    37  		{
    38  			"When given a non-keyframe; it should return false",
    39  			[]byte{0x27, 0x90, 0x90},
    40  			false,
    41  		},
    42  		{
    43  			"When given a SPS packetized with STAP-A; it should return true",
    44  			[]byte{0x38, 0x00, 0x03, 0x27, 0x90, 0x90, 0x00, 0x05, 0x28, 0x90, 0x90, 0x90, 0x90},
    45  			true,
    46  		},
    47  		{
    48  			"When given a SPS with no packetization; it should return true",
    49  			[]byte{0x27, 0x90, 0x90, 0x00},
    50  			true,
    51  		},
    52  	}
    53  
    54  	for _, tt := range tests {
    55  		tt := tt
    56  		t.Run(tt.name, func(t *testing.T) {
    57  			got := isKeyFrame(tt.payload)
    58  			assert.Equal(t, tt.want, got)
    59  		})
    60  	}
    61  }
    62  
    63  func TestWriteRTP(t *testing.T) {
    64  	tests := []struct {
    65  		name        string
    66  		payload     []byte
    67  		hasKeyFrame bool
    68  		wantBytes   []byte
    69  		wantErr     error
    70  		reuseWriter bool
    71  	}{
    72  		{
    73  			"When given an empty payload; it should return nil",
    74  			[]byte{},
    75  			false,
    76  			[]byte{},
    77  			nil,
    78  			false,
    79  		},
    80  		{
    81  			"When no keyframe is defined; it should discard the packet",
    82  			[]byte{0x25, 0x90, 0x90},
    83  			false,
    84  			[]byte{},
    85  			nil,
    86  			false,
    87  		},
    88  		{
    89  			"When a valid Single NAL Unit packet is given; it should unpack it without error",
    90  			[]byte{0x27, 0x90, 0x90},
    91  			true,
    92  			[]byte{0x00, 0x00, 0x00, 0x01, 0x27, 0x90, 0x90},
    93  			nil,
    94  			false,
    95  		},
    96  		{
    97  			"When a valid STAP-A packet is given; it should unpack it without error",
    98  			[]byte{0x38, 0x00, 0x03, 0x27, 0x90, 0x90, 0x00, 0x05, 0x28, 0x90, 0x90, 0x90, 0x90},
    99  			true,
   100  			[]byte{0x00, 0x00, 0x00, 0x01, 0x27, 0x90, 0x90, 0x00, 0x00, 0x00, 0x01, 0x28, 0x90, 0x90, 0x90, 0x90},
   101  			nil,
   102  			false,
   103  		},
   104  		{
   105  			"When a valid FU-A start packet is given; it should unpack it without error",
   106  			[]byte{0x3C, 0x85, 0x90, 0x90, 0x90},
   107  			true,
   108  			[]byte{},
   109  			nil,
   110  			true,
   111  		},
   112  		{
   113  			"When a valid FU-A end packet is given; it should unpack it without error",
   114  			[]byte{0x3C, 0x45, 0x90, 0x90, 0x90},
   115  			true,
   116  			[]byte{0x00, 0x00, 0x00, 0x01, 0x25, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90},
   117  			nil,
   118  			false,
   119  		},
   120  	}
   121  
   122  	var reuseWriter *bytes.Buffer
   123  	var reuseH264Writer *H264Writer
   124  
   125  	for _, tt := range tests {
   126  		tt := tt
   127  		t.Run(tt.name, func(t *testing.T) {
   128  			writer := &bytes.Buffer{}
   129  			h264Writer := &H264Writer{
   130  				hasKeyFrame: tt.hasKeyFrame,
   131  				writer:      writer,
   132  			}
   133  			if reuseWriter != nil {
   134  				writer = reuseWriter
   135  			}
   136  			if reuseH264Writer != nil {
   137  				h264Writer = reuseH264Writer
   138  			}
   139  
   140  			assert.Equal(t, tt.wantErr, h264Writer.WriteRTP(&rtp.Packet{
   141  				Payload: tt.payload,
   142  			}))
   143  			assert.True(t, bytes.Equal(tt.wantBytes, writer.Bytes()))
   144  
   145  			if !tt.reuseWriter {
   146  				assert.Nil(t, h264Writer.Close())
   147  				reuseWriter = nil
   148  				reuseH264Writer = nil
   149  			} else {
   150  				reuseWriter = writer
   151  				reuseH264Writer = h264Writer
   152  			}
   153  		})
   154  	}
   155  }