av/rtmp/rtmp_test.go

/*
NAME
  rtmp_test.go

DESCRIPTION
  RTMP tests

AUTHORS
  Saxon Nelson-Milton <saxon@ausocean.org>
  Dan Kortschak <dan@ausocean.org>
  Alan Noble <alan@ausocean.org>

LICENSE
  rtmp_test.go is Copyright (C) 2017-2019 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
*/

package rtmp

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"testing"
	"time"

	"bitbucket.org/ausocean/av/stream/flv"
	"bitbucket.org/ausocean/av/stream/lex"
)

const (
	rtmpProtocol = "rtmp"
	testHost     = "a.rtmp.youtube.com"
	testApp      = "live2"
	testBaseURL  = rtmpProtocol + "://" + testHost + "/" + testApp + "/"
	testTimeout  = 30
	testDataDir  = "../../test/test-data/av/input"
)

// testVerbosity controls the amount of output.
// NB: This is not the log level, which is DebugLevel.
//   0: suppress logging completely
//   1: log messages only
//   2: log messages with errors, if any
var testVerbosity = 1

// testKey is the YouTube RTMP key required for YouTube streaming (RTMP_TEST_KEY env var).
// NB: don't share your key with others.
var testKey string

// testFile is the test video file  (RTMP_TEST_FILE env var).
// betterInput.h264 is a good one to use.
var testFile string

// testLog is a bare bones logger that logs to stdout, and exits upon either an error or fatal error.
func testLog(level int8, msg string, params ...interface{}) {
	logLevels := [...]string{"Debug", "Info", "Warn", "Error", "", "", "Fatal"}
	if testVerbosity == 0 {
		return
	}
	if level < -1 || level > 5 {
		panic("Invalid log level")
	}
	switch testVerbosity {
	case 0:
		// silence is golden
	case 1:
		fmt.Printf("%s: %s\n", logLevels[level+1], msg)
	case 2:
		// extract the first param if it is one we care about, otherwise just print the message
		if len(params) >= 2 {
			switch params[0].(string) {
			case "error":
				fmt.Printf("%s: %s, error=%v\n", logLevels[level+1], msg, params[1].(string))
			case "size":
				fmt.Printf("%s: %s, size=%d\n", logLevels[level+1], msg, params[1].(int))
			default:
				fmt.Printf("%s: %s\n", logLevels[level+1], msg)
			}
		} else {
			fmt.Printf("%s: %s\n", logLevels[level+1], msg)
		}
	}
	if level >= 4 {
		// Error or Fatal
		buf := make([]byte, 1<<16)
		size := runtime.Stack(buf, true)
		fmt.Printf("%s\n", string(buf[:size]))
		os.Exit(1)
	}
}

// TestKey tests that the RTMP_TEST_KEY environment variable is present
func TestKey(t *testing.T) {
	testLog(0, "TestKey")
	testKey = os.Getenv("RTMP_TEST_KEY")
	if testKey == "" {
		msg := "RTMP_TEST_KEY environment variable not defined"
		testLog(0, msg)
		t.Skip(msg)
	}
	testLog(0, "Testing against URL "+testBaseURL+testKey)
}

// TestSetupURL tests URL parsing.
func TestSetupURL(t *testing.T) {
	testLog(0, "TestSetupURL")
	// test with just the base URL
	s := NewSession(testBaseURL, testTimeout, testLog)
	if s.url != testBaseURL && s.link.timeout != testTimeout {
		t.Errorf("NewSession failed")
	}
	err := setupURL(s)
	if err != nil {
		t.Errorf("setupURL(testBaseURL) failed with error: %v", err)
	}
	// test the parts are as expected
	if rtmpProtocolStrings[s.link.protocol] != rtmpProtocol {
		t.Errorf("setupURL returned wrong protocol: %v", s.link.protocol)
	}
	if s.link.host != testHost {
		t.Errorf("setupURL returned wrong host: %v", s.link.host)
	}
	if s.link.app != testApp {
		t.Errorf("setupURL returned wrong app: %v", s.link.app)
	}
}

// TestOpenClose tests opening an closing an RTMP connection.
func TestOpenClose(t *testing.T) {
	testLog(0, "TestOpenClose")
	if testKey == "" {
		t.Skip("Skipping TestOpenClose since no RTMP_TEST_KEY")
	}
	s := NewSession(testBaseURL+testKey, testTimeout, testLog)
	err := s.Open()
	if err != nil {
		t.Errorf("Open failed with error: %v", err)
		return
	}
	err = s.Close()
	if err != nil {
		t.Errorf("Close failed with error: %v", err)
		return
	}
}

// TestFromFrame tests streaming from a single H.264 frame which is repeated.
func TestFromFrame(t *testing.T) {
	testLog(0, "TestFromFrame")
	if testKey == "" {
		t.Skip("Skipping TestFromFrame since no RTMP_TEST_KEY")
	}
	s := NewSession(testBaseURL+testKey, testTimeout, testLog)
	err := s.Open()
	if err != nil {
		t.Errorf("Session.Open failed with error: %v", err)
	}

	// TODO: don't hard code this path
	b, err := ioutil.ReadFile("/home/saxon/Desktop/singleframe.h264")
	if err != nil {
		t.Errorf("ReadFile failed with error: %v", err)
	}

	const noOfFrames = 1000
	videoData := make([]byte, 0, noOfFrames*len(b))
	for i := 0; i < noOfFrames; i++ {
		videoData = append(videoData, b...)
	}

	const frameRate = 25
	flvEncoder := flv.NewEncoder(s, true, true, frameRate)
	err = lex.H264(flvEncoder, bytes.NewReader(videoData), time.Second/time.Duration(frameRate))
	if err != nil {
		t.Errorf("Lexing failed with error: %v", err)
	}
	err = s.Close()
	if err != nil {
		t.Errorf("Session.Close failed with error: %v", err)
	}
}

// TestFromFile tests streaming from an video file comprising raw H.264.
// The test file is supplied via the RTMP_TEST_FILE environment variable.
func TestFromFile(t *testing.T) {
	testLog(0, "TestFromFile")
	testFile := os.Getenv("RTMP_TEST_FILE")
	if testFile == "" {
		t.Skip("Skipping TestFromFile since no RTMP_TEST_FILE")
	}
	if testKey == "" {
		t.Skip("Skipping TestFromFile since no RTMP_TEST_KEY")
	}
	s := NewSession(testBaseURL+testKey, testTimeout, testLog)
	err := s.Open()
	if err != nil {
		t.Errorf("Session.Open failed with error: %v", err)
	}

	f, err := os.Open(testFile)
	if err != nil {
		t.Errorf("Open failed with error: %v", err)
	}
	defer f.Close()

	// Pass RTMP session, true for audio, true for video, and 25 FPS
	flvEncoder := flv.NewEncoder(s, true, true, 25)
	err = lex.H264(flvEncoder, f, time.Second/time.Duration(25))
	if err != nil {
		t.Errorf("Lexing and encoding failed with error: %v", err)
	}

	err = s.Close()
	if err != nil {
		t.Errorf("Session.Close failed with error: %v", err)
	}
}