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"
	"os"
	"runtime"
	"testing"
	"io/ioutil"
)

const (
	rtmpProtocol = "rtmp"
	testHost     = "a.rtmp.youtube.com"
	testApp      = "live2"
	testBaseURL  = rtmpProtocol + "://" + testHost + "/" + testApp + "/"
	testTimeout  = 30
)

// debug enables extra logging
var testDebug bool

// 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 level < -1 || level > 5 {
		panic("Invalid log level")
	}
	if testDebug && 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 == 5 {
		// Fatal
		buf := make([]byte, 1<<16)
		size := runtime.Stack(buf, true)
		fmt.Printf("%s\n", string(buf[:size]))
		os.Exit(1)
	}
}

func TestKey(t *testing.T) {
	testLog(0, "TestKey")
	testKey = os.Getenv("RTMP_TEST_KEY")
	if testKey == "" {
		fmt.Printf("RTMP_TEST_KEY environment variable not defined\n")
		os.Exit(1)
	}
	testLog(0, "Testing against URL "+testBaseURL+testKey)
}

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 again with the full URL
	s = NewSession(testBaseURL+testKey, testTimeout, testLog)
	err = setupURL(s, s.url)
	if err != nil {
		t.Errorf("setupURL(testBaseURL+testKey) 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)
	}
}

func TestOpen(t *testing.T) {
	testLog(0, "TestOpen")
	s := NewSession(testBaseURL+testKey, testTimeout, testLog)
	err := setupURL(s, s.url)
	if err != nil {
		t.Errorf("setupURL failed with error: %v", err)
	}
	s.enableWrite()
	err = s.Open()
	if err != nil {
		t.Errorf("connect failed with error: %v", err)
	}
}

func TestOpenClose(t *testing.T) {
	testLog(0, "TestOpenClose")
	s := NewSession(testBaseURL+testKey, testTimeout, testLog)
	err := s.Open()
	if err != nil {
		t.Errorf("Open failed with error: %v", err)
		return
	}
}

func TestFromFile(t *testing.T) {
	testLog(0, "TestFromFile")
	testFile := os.Getenv("RTMP_TEST_FILE")
	if testKey == "" {
		fmt.Printf("RTMP_TEST_FILE environment variable not defined\n")
		os.Exit(1)
	}
	s := NewSession(testBaseURL+testKey, testTimeout, testLog)
	err := s.Open()
	if err != nil {
		t.Errorf("Session.Open failed with error: %v", err)
	}
	video, err := ioutil.ReadFile(testFile)
	if err != nil {
		t.Errorf("Cannot open video file: %v", testFile)
	}
	// ToDo: rate limit writing
	n, err := s.Write(video)

	if err != nil {
		t.Errorf("Session.Write failed with error: %v", err)
	}
	if n != len(video) {
		t.Errorf("Session.Write retuned wrong length")
	}
	err = s.Close()
	if err != nil {
		t.Errorf("Session.Close failed with error: %v", err)
	}
}