From 72cd0b7a034c97540cdbf8dfaadb42b4f946c0f7 Mon Sep 17 00:00:00 2001
From: saxon <saxon.milton@gmail.com>
Date: Wed, 4 Jul 2018 14:57:23 +0930
Subject: [PATCH] Got simple frame check working

---
 parser/h264.go        |  2 +-
 parser/parser_test.go | 37 +++++++++++++++++++------------------
 2 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/parser/h264.go b/parser/h264.go
index 630c6ee2..86ecd030 100644
--- a/parser/h264.go
+++ b/parser/h264.go
@@ -130,7 +130,7 @@ func (p *H264) parse() {
 					return
 				}
 				outputBuffer = append(outputBuffer, aByte)
-				if nalType := aByte & 0x1F; nalType == 1 || nalType == 5 || nalType == 8 {
+				if nalType := aByte & 0x1F; nalType == 1 || nalType == 5 || nalType == 8 || nalType == 7 {
 					searchingForEnd = true
 				}
 			}
diff --git a/parser/parser_test.go b/parser/parser_test.go
index f79454ec..94be834e 100644
--- a/parser/parser_test.go
+++ b/parser/parser_test.go
@@ -9,41 +9,52 @@ import (
 
 const (
 	mjpegInputFileName = "testInput/testInput.avi"
-	h264InputFileName = "bitbucket.org/ausocean/test/test-data/av/input/betterInput.h264"
+	h264InputFileName = "../../test/test-data/av/input/betterInput.h264"
 )
 
 func TestH264Parser(t *testing.T) {
 	fmt.Println("Opening input file!")
-	// Open the input file
 	inputFile, err := os.Open(h264InputFileName)
 	if err != nil {
 		t.Errorf("Should not have got error opening file!")
+		return
 	}
+
 	fmt.Println("Getting file stats!")
 	stats, err := inputFile.Stat()
 	if err != nil {
 		t.Errorf("Could not get input file stats!")
 		return
 	}
-	fmt.Println("Creating space for file data!")
+
+	fmt.Println("Creating space for file data and reading!")
 	data := make([]byte, stats.Size())
 	_, err = inputFile.Read(data)
 	if err != nil {
 		t.Errorf("Should not have got read error!")
 		return
 	}
+
+	// Create 'parser' and start it up
 	fmt.Println("Creating parser!")
-	parser := NewH264Parser(len(data) + 1)
+	parser := NewH264Parser()
 	parser.SetOutputChan(make(chan []byte, 10000))
 	parser.Start()
-	fmt.Printf("len(data): %v\n", len(data))
 
 	count := 0
 	for i := range data {
-		parser.GetInputChan() <- data[i]
+		// If we have a whole frame, write it to a file for inspection otherwise do
+		// nothing
 		select {
-		case frame:=<-parser.GetOutputChan()
-		outputFile, err := os.Create("")
+		case parser.InputChan() <- data[i]:
+		case frame:=<-parser.OutputChan():
+			outputFile, err := os.Create("testOutput/" + strconv.Itoa(count) + ".h264_frame")
+			if err != nil {
+				t.Errorf("Should not have got error creating output file!")
+				return
+			}
+			outputFile.Write(frame)
+			outputFile.Close()
 			count++
 			if count > 4 {
 				return
@@ -51,16 +62,6 @@ func TestH264Parser(t *testing.T) {
 		default:
 		}
 	}
-	for i := 0; len(parser.GetOutputChan()) > 0; i++ {
-		// Open a new output file
-		outputFile, err := os.Create("testOutput/image" + strconv.Itoa(i) + ".jpeg")
-		if err != nil {
-			t.Errorf("Should not have got error creating output file!")
-			return
-		}
-		outputFile.Write(<-parser.GetOutputChan())
-		outputFile.Close()
-	}
 }