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"
	"time"

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

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

// debug enables extra logging
var testDebug = true

// 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")
	s := NewSession(testBaseURL+testKey, testTimeout, testLog)
	err := s.Open()
	if err != nil {
		t.Errorf("Session.Open failed with error: %v", err)
	}

	// This read from a h264 file
	if true {
		// Open file
		f, err := os.Open("../../test/test-data/av/input/betterInput.h264")
		if err != nil {
			t.Errorf("Cannot open video file: %v", testFile)
		}
		defer f.Close()

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

		// This reads a single h264 frame and sends
	} else {
		b, err := ioutil.ReadFile("/home/saxon/Downloads/ausoceanFrame.h264") // b has type []byte
		if err != nil {
			t.Errorf("Could not read file, failed with error: %v", err)
		}
		flvEncoder := flv.NewEncoder(s, true, true, 25)
		for i := 0; i < 10000; i++ {
			err := flvEncoder.Encode(b)
			if err != nil {
				t.Errorf("Encoding failed!")
			}
		}
	}

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