av/rtmp/rtmp_test.go

/*
NAME
  rtmp_test.go

DESCRIPTION
  RTMP tests

AUTHORS
  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 (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"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 RTMP key required for YouTube streaming (RTMP_TEST_KEY env var)
var testKey string

// testFile is the test video file  (RTMP_TEST_FILE env var)
var testFile string

// testLog is a bare bones logger that logs to stdout.
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")
	}
	if testVerbosity == 2 && len(params) >= 2 {
		// extract the params we know about, otherwise just print the message
		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 == 5 {
		// 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, s.url)
	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)
	}

	b, err := ioutil.ReadFile(filepath.Join(testDataDir, "AusOcean_logo_1080p.h264"))
	if err != nil {
		t.Errorf("ReadFile failed with error: %v", err)
	}

	// Pass RTMP session, true for audio, true for video, and 25 FPS
	flvEncoder := flv.NewEncoder(s, true, true, 25)
	for i := 0; i < 25; i++ {
		err := flvEncoder.Encode(b)
		if err != nil {
			t.Errorf("Encoding failed with error: %v", err)
		}
		time.Sleep(40 * time.Millisecond) // rate limit to 1/25s
	}

	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)
	}
}