clipper/backend/media/service_test.go

package media_test

import (
	"bytes"
	"context"
	"database/sql"
	"io"
	"os"
	"testing"

	"git.netflux.io/rob/clipper/config"
	"git.netflux.io/rob/clipper/generated/mocks"
	"git.netflux.io/rob/clipper/generated/store"
	"git.netflux.io/rob/clipper/media"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
)

func TestPeaksForSegment(t *testing.T) {
	testCases := []struct {
		name                 string
		fixturePath          string
		fixtureLen           int64
		startFrame, endFrame int64
		channels             int32
		numBins              int
		wantPeaks            []int16
		wantErr              string
	}{
		{
			name:        "entire fixture, stereo, 1 bin",
			fixturePath: "testdata/tone-44100-stereo-int16.raw",
			fixtureLen:  176400,
			startFrame:  0,
			endFrame:    44100,
			channels:    2,
			numBins:     1,
			wantPeaks:   []int16{32747, 32747},
		},
		{
			name:        "entire fixture, stereo, 4 bins",
			fixturePath: "testdata/tone-44100-stereo-int16.raw",
			fixtureLen:  176400,
			startFrame:  0,
			endFrame:    44100,
			channels:    2,
			numBins:     4,
			wantPeaks:   []int16{8173, 8177, 16366, 16370, 24557, 24555, 32747, 32747},
		},
		{
			name:        "entire fixture, stereo, 16 bins",
			fixturePath: "testdata/tone-44100-stereo-int16.raw",
			fixtureLen:  176400,
			startFrame:  0,
			endFrame:    44100,
			channels:    2,
			numBins:     16,
			wantPeaks:   []int16{2029, 2029, 4075, 4076, 6124, 6125, 8173, 8177, 10222, 10221, 12267, 12265, 14314, 14313, 16366, 16370, 18413, 18411, 20453, 20454, 22505, 22508, 24557, 24555, 26604, 26605, 28644, 28643, 30698, 30694, 32747, 32747},
		},
		{
			name:        "entire fixture, mono, 1 bin",
			fixturePath: "testdata/tone-44100-mono-int16.raw",
			fixtureLen:  88200,
			startFrame:  0,
			endFrame:    44100,
			channels:    1,
			numBins:     1,
			wantPeaks:   []int16{32748},
		},
		{
			name:        "entire fixture, mono, 32 bins",
			fixturePath: "testdata/tone-44100-mono-int16.raw",
			fixtureLen:  88200,
			startFrame:  0,
			endFrame:    44100,
			channels:    1,
			numBins:     32,
			wantPeaks:   []int16{1026, 2030, 3071, 4075, 5122, 6126, 7167, 8172, 9213, 10217, 11259, 12264, 13311, 14315, 15360, 16364, 17405, 18412, 19450, 20453, 21497, 22504, 23549, 24554, 25599, 26607, 27641, 28642, 29688, 30738, 31746, 32748},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			startByte := tc.startFrame * int64(tc.channels) * media.SizeOfInt16
			endByte := tc.endFrame * int64(tc.channels) * media.SizeOfInt16
			expectedBytes := endByte - startByte

			audioFile, err := os.Open(tc.fixturePath)
			require.NoError(t, err)
			defer audioFile.Close()
			audioData := io.NopCloser(io.LimitReader(audioFile, int64(expectedBytes)))

			mediaSet := store.MediaSet{
				ID:            uuid.New(),
				AudioChannels: tc.channels,
				AudioRawS3Key: sql.NullString{String: "foo", Valid: true},
			}

			// store is passed the mediaSetID and returns a mediaSet
			store := &mocks.Store{}
			store.On("GetMediaSet", mock.Anything, mediaSet.ID).Return(mediaSet, nil)
			defer store.AssertExpectations(t)

			// fileStore is passed the expected byte range, and returns an io.Reader
			fileStore := &mocks.FileStore{}
			fileStore.
				On("GetObjectWithRange", mock.Anything, "foo", startByte, endByte).
				Return(audioData, nil)

			service := media.NewMediaSetService(store, nil, fileStore, nil, media.NewTestWorkerPool(), config.Config{}, zap.NewNop().Sugar())
			peaks, err := service.GetPeaksForSegment(context.Background(), mediaSet.ID, tc.startFrame, tc.endFrame, tc.numBins)

			if tc.wantErr == "" {
				assert.NoError(t, err)
				assert.Equal(t, tc.wantPeaks, peaks)
			} else {
				assert.EqualError(t, err, tc.wantErr)
			}
		})
	}
}

func BenchmarkGetPeaksForSegment(b *testing.B) {
	const (
		startFrame  = 0
		endFrame    = 1323000
		channels    = 2
		fixturePath = "testdata/tone-44100-stereo-int16-30000ms.raw"
		fixtureLen  = 5292000
		numBins     = 2000
	)

	audioFile, err := os.Open(fixturePath)
	require.NoError(b, err)
	audioData, err := io.ReadAll(audioFile)
	require.NoError(b, err)

	mediaSetID := uuid.New()
	mediaSet := store.MediaSet{ID: mediaSetID, AudioChannels: channels}

	store := &mocks.Store{}
	store.On("GetMediaSet", mock.Anything, mediaSetID).Return(mediaSet, nil)

	for n := 0; n < b.N; n++ {
		// recreate the reader on each iteration
		b.StopTimer()
		readCloser := io.NopCloser(bytes.NewReader(audioData))
		fileStore := &mocks.FileStore{}
		fileStore.
			On("GetObjectWithRange", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
			Return(readCloser, nil)

		service := media.NewMediaSetService(store, nil, fileStore, nil, media.NewTestWorkerPool(), config.Config{}, zap.NewNop().Sugar())
		b.StartTimer()

		_, err = service.GetPeaksForSegment(context.Background(), mediaSetID, startFrame, endFrame, numBins)
		require.NoError(b, err)
	}
}