Added media streaming tests and improved comments.

rtmp/rtmp_test.go

@@ -29,14 +29,15 @@ package rtmp
 
 import (
 	"fmt"
+	"io/ioutil"
 	"os"
+	"path/filepath"
 	"runtime"
 	"testing"
 	"time"
 
 	"bitbucket.org/ausocean/av/stream/flv"
 	"bitbucket.org/ausocean/av/stream/lex"
-	"io/ioutil"
 )
 
 const (
@@ -45,10 +46,15 @@ const (
 	testApp      = "live2"
 	testBaseURL  = rtmpProtocol + "://" + testHost + "/" + testApp + "/"
 	testTimeout  = 30
+	testDataDir  = "../../test/test-data/av/input"
 )
 
-// debug enables extra logging
+// testVerbosity controls the amount of output
-var testDebug = true
+// 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
@@ -59,10 +65,14 @@ 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 testDebug && len(params) >= 2 {
+	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))
@@ -83,16 +93,19 @@ func testLog(level int8, msg string, params ...interface{}) {
 	}
 }
 
+// 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 == "" {
-		fmt.Printf("RTMP_TEST_KEY environment variable not defined\n")
+		msg := "RTMP_TEST_KEY environment variable not defined"
-		os.Exit(1)
+		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
@@ -104,12 +117,6 @@ func TestSetupURL(t *testing.T) {
 	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)
@@ -122,66 +129,48 @@ func TestSetupURL(t *testing.T) {
 	}
 }
 
-func TestOpen(t *testing.T) {
+// TestOpenClose tests opening an closing an RTMP connection.
-	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")
+	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
+	}
 }
 
-func TestFromFile(t *testing.T) {
+// TestFromFrame tests streaming from a single H.264 frame which is repeated.
-	testLog(0, "TestFromFile")
+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)
 	}
 
-	// This read from a h264 file
+	b, err := ioutil.ReadFile(filepath.Join(testDataDir, "AusOcean_logo_1080p.h264"))
-	if true {
+	if err != nil {
-		// Open file
+		t.Errorf("ReadFile failed with error: %v", err)
-		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
+	// Pass RTMP session, true for audio, true for video, and 25 FPS
-		flvEncoder := flv.NewEncoder(s, true, true, 25)
+	flvEncoder := flv.NewEncoder(s, true, true, 25)
-		err = lex.H264(flvEncoder, f, time.Second/time.Duration(25))
+	for i := 0; i < 25; i++ {
+		err := flvEncoder.Encode(b)
 		if err != nil {
-			t.Errorf("Lexing and encoding failed with error: %v", err)
+			t.Errorf("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!")
-			}
 		}
 	}
 
@@ -190,3 +179,39 @@ func TestFromFile(t *testing.T) {
 		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)
+	}
+}