av/revid/revid_test.go

package revid

import (
	"errors"
	"fmt"
	"io"
	"os"
	"runtime"
	"testing"

	"bitbucket.org/ausocean/iot/pi/netsender"
)

const raspividPath = "/usr/local/bin/raspivid"

// Suppress all test logging, except for t.Errorf output.
var silent bool

// TestRaspivid tests that raspivid starts correctly.
// It is intended to be run on a Raspberry Pi.
func TestRaspivid(t *testing.T) {
	if _, err := os.Stat(raspividPath); os.IsNotExist(err) {
		t.Skip("Skipping TestRaspivid since no raspivid found.")
	}

	var logger testLogger
	ns, err := netsender.New(&logger, nil, nil, nil)
	if err != nil {
		t.Errorf("netsender.New failed with error %v", err)
	}

	var c Config
	c.Logger = &logger
	c.Input = Raspivid
	c.Outputs = make([]uint8, 1)

	rv, err := New(c, ns)
	if err != nil {
		t.Errorf("revid.New failed with error %v", err)
	}

	err = rv.Start()
	if err != nil {
		t.Errorf("revid.Start failed with error %v", err)
	}
}

// testLogger implements a netsender.Logger.
type testLogger struct{}

// SetLevel normally sets the logging level, but it is a no-op in our case.
func (tl *testLogger) SetLevel(level int8) {
}

// Log requests the Logger to write a message at the given level.
func (tl *testLogger) Log(level int8, msg string, params ...interface{}) {
	logLevels := [...]string{"Debug", "Info", "Warn", "Error", "", "", "Fatal"}
	if level < -1 || level > 5 {
		panic("Invalid log level")
	}
	if !silent {
		fmt.Printf("%s: %s\n", logLevels[level+1], msg)
	}
	if level == 5 {
		buf := make([]byte, 1<<16)
		size := runtime.Stack(buf, true)
		fmt.Printf("%s\n", string(buf[:size]))
		os.Exit(1)
	}
}

// tstMtsEncoder emulates the mts.Encoder to the extent of the dst field.
// This will allow access to the dst to check that it has been set corrctly.
type tstMtsEncoder struct {
	dst io.Writer
}

// newTstMtsEncoder returns a pointer to a newTsMtsEncoder.
func newTstMtsEncoder(dst io.Writer, fps int) io.Writer {
	return &tstMtsEncoder{dst: dst}
}

func (e *tstMtsEncoder) Write(d []byte) (int, error) { return 0, nil }

// tstFlvEncoder emulates the flv.Encoder to the extent of the dst field.
// This will allow access to the dst to check that it has been set corrctly.
type tstFlvEncoder struct {
	dst io.Writer
}

// newTstFlvEncoder returns a pointer to a new tstFlvEncoder.
func newTstFlvEncoder(dst io.Writer, fps int) (io.Writer, error) {
	return &tstFlvEncoder{dst: dst}, nil
}

func (e *tstFlvEncoder) Write(d []byte) (int, error) { return 0, nil }

// TestResetEncoderSenderSetup checks that revid.reset() correctly sets up the
// revid.encoder slice and the senders the encoders write to.
func TestResetEncoderSenderSetup(t *testing.T) {
	// We will use these to indicate types after assertion.
	const (
		mtsSenderStr  = "revid.mtsSender"
		rtpSenderStr  = "revid.rtpSender"
		rtmpSenderStr = "revid.RtmpSender"
		mtsEncoderStr = "mts.Encoder"
		flvEncoderStr = "flv.Encoder"
	)

	// Struct that will be used to format test cases nicely below.
	type encoder struct {
		encoderType  string
		destinations []string
	}

	tests := []struct {
		outputs  []uint8
		encoders []encoder
	}{
		{
			outputs: []uint8{Http},
			encoders: []encoder{
				{
					encoderType:  mtsEncoderStr,
					destinations: []string{mtsSenderStr},
				},
			},
		},
		{
			outputs: []uint8{Rtmp},
			encoders: []encoder{
				{
					encoderType:  flvEncoderStr,
					destinations: []string{rtmpSenderStr},
				},
			},
		},
		{
			outputs: []uint8{Rtp},
			encoders: []encoder{
				{
					encoderType:  mtsEncoderStr,
					destinations: []string{rtpSenderStr},
				},
			},
		},
		{
			outputs: []uint8{Http, Rtmp},
			encoders: []encoder{
				{
					encoderType:  mtsEncoderStr,
					destinations: []string{mtsSenderStr},
				},
				{
					encoderType:  flvEncoderStr,
					destinations: []string{rtmpSenderStr},
				},
			},
		},
		{
			outputs: []uint8{Http, Rtp, Rtmp},
			encoders: []encoder{
				{
					encoderType:  mtsEncoderStr,
					destinations: []string{mtsSenderStr, rtpSenderStr},
				},
				{
					encoderType:  flvEncoderStr,
					destinations: []string{rtmpSenderStr},
				},
			},
		},
		{
			outputs: []uint8{Rtp, Rtmp},
			encoders: []encoder{
				{
					encoderType:  mtsEncoderStr,
					destinations: []string{rtpSenderStr},
				},
				{
					encoderType:  flvEncoderStr,
					destinations: []string{rtmpSenderStr},
				},
			},
		},
	}

	// typeOfEncoder will return the type of encoder implementing stream.Encoder.
	typeOfEncoder := func(i io.Writer) (string, error) {
		if _, ok := i.(*tstMtsEncoder); ok {
			return mtsEncoderStr, nil
		}
		if _, ok := i.(*tstFlvEncoder); ok {
			return flvEncoderStr, nil
		}
		return "", errors.New("unknown Encoder type")
	}

	// typeOfSender will return the type of sender implementing loadSender.
	typeOfSender := func(s loadSender) (string, error) {
		if _, ok := s.(*mtsSender); ok {
			return mtsSenderStr, nil
		}
		if _, ok := s.(*rtpSender); ok {
			return rtpSenderStr, nil
		}
		if _, ok := s.(*rtmpSender); ok {
			return rtmpSenderStr, nil
		}
		return "", errors.New("unknown loadSender type")
	}

	rv, err := New(Config{Logger: &testLogger{}}, nil)
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	// Go through our test cases.
	for testNum, test := range tests {
		// Create a new config and reset revid with it.
		const dummyURL = "rtmp://dummy"
		c := Config{Logger: &testLogger{}, Outputs: test.outputs, RtmpUrl: dummyURL}
		err := rv.setConfig(c)
		if err != nil {
			t.Fatalf("unexpected error: %v for test %v", err, testNum)
		}

		// This logic is what we want to check.
		err = rv.setupPipeline(newTstMtsEncoder, newTstFlvEncoder)
		if err != nil {
			t.Fatalf("unexpected error: %v for test %v", err, testNum)
		}

		// First check that we have the correct number of encoders.
		got := len(rv.encoder)
		want := len(test.encoders)
		if got != want {
			t.Errorf("incorrect number of encoders in revid for test: %v. \nGot: %v\nWant: %v\n", testNum, got, want)
		}

		// Now check the correctness of encoders and their destinations.
		for _, e := range rv.encoder {
			// Get e's type.
			encoderType, err := typeOfEncoder(e)
			if err != nil {
				t.Fatalf("could not get encoders type for test %v, failed with err: %v", testNum, err)
			}

			// Check that we expect this encoder to be here.
			idx := -1
			for i, expect := range test.encoders {
				if expect.encoderType == encoderType {
					idx = i
				}
			}
			if idx == -1 {
				t.Fatalf("encoder %v isn't expected in test %v", encoderType, testNum)
			}

			// Now check that this encoder has correct number of destinations (senders).
			var ms io.Writer
			switch encoderType {
			case mtsEncoderStr:
				ms = e.(*tstMtsEncoder).dst
			case flvEncoderStr:
				ms = e.(*tstFlvEncoder).dst
			}

			senders := ms.(*multiSender).senders
			got = len(senders)
			want = len(test.encoders[idx].destinations)
			if got != want {
				t.Errorf("did not get expected number of senders in test %v. \nGot: %v\nWant: %v\n", testNum, got, want)
			}

			// Check that destinations are as expected.
			for _, expectDst := range test.encoders[idx].destinations {
				ok := false
				for _, dst := range senders {
					// Get type of sender.
					senderType, err := typeOfSender(dst)
					if err != nil {
						t.Fatalf("could not get encoders type for test %v, failed with err: %v", testNum, err)
					}

					// If it's one we want, indicate.
					if senderType == expectDst {
						ok = true
					}
				}

				// If not okay then we couldn't find expected sender.
				if !ok {
					t.Errorf("could not find expected destination %v, for test %v", expectDst, testNum)
				}
			}
		}
	}
}