Wait for highlighting to complete in pager tests

m/pager_test.go

@@ -13,9 +13,6 @@ import (
 
 func TestUnicodeRendering(t *testing.T) {
 	reader := NewReaderFromStream("", strings.NewReader("åäö"))
-	if err := reader._Wait(); err != nil {
-		panic(err)
-	}
 
 	var answers = []twin.Cell{
 		twin.NewCell('å', twin.StyleDefault),
@@ -40,9 +37,6 @@ func logDifference(t *testing.T, expected twin.Cell, actual twin.Cell) {
 func TestFgColorRendering(t *testing.T) {
 	reader := NewReaderFromStream("", strings.NewReader(
 		"\x1b[30ma\x1b[31mb\x1b[32mc\x1b[33md\x1b[34me\x1b[35mf\x1b[36mg\x1b[37mh\x1b[0mi"))
-	if err := reader._Wait(); err != nil {
-		panic(err)
-	}
 
 	var answers = []twin.Cell{
 		twin.NewCell('a', twin.StyleDefault.Foreground(twin.NewColor16(0))),
@@ -65,9 +59,6 @@ func TestFgColorRendering(t *testing.T) {
 func TestBrokenUtf8(t *testing.T) {
 	// The broken UTF8 character in the middle is based on "©" = 0xc2a9
 	reader := NewReaderFromStream("", strings.NewReader("abc\xc2def"))
-	if err := reader._Wait(); err != nil {
-		panic(err)
-	}
 
 	var answers = []twin.Cell{
 		twin.NewCell('a', twin.StyleDefault),
@@ -86,6 +77,11 @@ func TestBrokenUtf8(t *testing.T) {
 }
 
 func startPaging(t *testing.T, reader *Reader) *twin.FakeScreen {
+	err := reader._Wait()
+	if err != nil {
+		panic(err)
+	}
+
 	screen := twin.NewFakeScreen(20, 10)
 	pager := NewPager(reader)
 	pager.ShowLineNumbers = false
@@ -105,9 +101,6 @@ func startPaging(t *testing.T, reader *Reader) *twin.FakeScreen {
 // assertIndexOfFirstX verifies the (zero-based) index of the first 'x'
 func assertIndexOfFirstX(t *testing.T, s string, expectedIndex int) {
 	reader := NewReaderFromStream("", strings.NewReader(s))
-	if err := reader._Wait(); err != nil {
-		panic(err)
-	}
 
 	contents := startPaging(t, reader).GetRow(0)
 	for pos, cell := range contents {
@@ -158,9 +151,6 @@ func TestCodeHighlighting(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
-	if err := reader._Wait(); err != nil {
-		panic(err)
-	}
 
 	packageKeywordStyle := twin.StyleDefault.WithAttr(twin.AttrBold).Foreground(twin.NewColorHex(0x6AB825))
 	packageNameStyle := twin.StyleDefault.Foreground(twin.NewColorHex(0xD0D0D0))
@@ -184,9 +174,6 @@ func TestCodeHighlighting(t *testing.T) {
 
 func testManPageFormatting(t *testing.T, input string, expected twin.Cell) {
 	reader := NewReaderFromStream("", strings.NewReader(input))
-	if err := reader._Wait(); err != nil {
-		panic(err)
-	}
 
 	// Without these three lines the man page tests will fail if either of these
 	// environment variables are set when the tests are run.

m/reader.go

@@ -32,8 +32,9 @@ type Reader struct {
 	// Have we had our contents replaced using setText()?
 	replaced bool
 
-	done           chan bool
-	moreLinesAdded chan bool
+	done             chan bool
+	highlightingDone chan bool // Used by tests
+	moreLinesAdded   chan bool
 }
 
 // Lines contains a number of lines from the reader, plus metadata
@@ -166,6 +167,7 @@ func readStream(stream io.Reader, reader *Reader, fromFilter *exec.Cmd) {
 // corner to help the user keep track of what is being paged.
 func NewReaderFromStream(name string, reader io.Reader) *Reader {
 	mReader := newReaderFromStream(reader, nil)
+	mReader.highlightingDone <- true // No highlighting of streams = nothing left to do = Done!
 
 	if len(name) > 0 {
 		mReader.lock.Lock()
@@ -184,6 +186,7 @@ func newReaderFromStream(reader io.Reader, fromFilter *exec.Cmd) *Reader {
 	var lines []*Line
 	var lock = &sync.Mutex{}
 	done := make(chan bool, 1)
+	highlightingDone := make(chan bool, 1)
 
 	// This needs to be size 1. If it would be 0, and we add more lines while the
 	// pager is processing, the pager would miss the lines added while it was
@@ -191,10 +194,11 @@ func newReaderFromStream(reader io.Reader, fromFilter *exec.Cmd) *Reader {
 	moreLinesAdded := make(chan bool, 1)
 
 	returnMe := Reader{
-		lines:          lines,
-		lock:           lock,
-		done:           done,
-		moreLinesAdded: moreLinesAdded,
+		lines:            lines,
+		lock:             lock,
+		done:             done,
+		moreLinesAdded:   moreLinesAdded,
+		highlightingDone: highlightingDone,
 	}
 
 	// FIXME: Make sure that if we panic somewhere inside of this goroutine,
@@ -217,19 +221,26 @@ func NewReaderFromText(name string, text string) *Reader {
 		}
 	}
 	done := make(chan bool, 1)
+	highlightingDone := make(chan bool, 1)
 	done <- true
-	return &Reader{
-		name:  &name,
-		lines: lines,
-		lock:  &sync.Mutex{},
-		done:  done,
+	returnMe := &Reader{
+		name:             &name,
+		lines:            lines,
+		lock:             &sync.Mutex{},
+		done:             done,
+		highlightingDone: highlightingDone,
 	}
+
+	returnMe.highlightingDone <- true // No highlighting to do = nothing left = Done!
+
+	return returnMe
 }
 
-// Wait for reader to finish reading. Used by tests.
+// Wait for reader to finish reading and highlighting. Used by tests.
 func (r *Reader) _Wait() error {
 	// Wait for our goroutine to finish
 	<-r.done
+	<-r.highlightingDone
 
 	r.lock.Lock()
 	defer r.lock.Unlock()
@@ -259,6 +270,7 @@ func newReaderFromCommand(filename string, filterCommand ...string) (*Reader, er
 	}
 
 	reader := newReaderFromStream(filterOut, filter)
+	reader.highlightingDone <- true // No highlighting to do == nothing left == Done!
 	reader.lock.Lock()
 	reader.name = &filename
 	reader._stderr = filterErr
@@ -316,14 +328,17 @@ func NewReaderFromFilename(filename string) (*Reader, error) {
 	if err != nil {
 		return nil, err
 	}
-	returnMe := NewReaderFromStream(filename, stream)
 
-	if strings.HasSuffix(filename, ".txt") {
-		// Not much point in highlighting text files, prefer streaming these
-		return returnMe, nil
-	}
+	returnMe := newReaderFromStream(stream, nil)
+	returnMe.lock.Lock()
+	returnMe.name = &filename
+	returnMe.lock.Unlock()
 
 	go func() {
+		defer func() {
+			returnMe.highlightingDone <- true
+		}()
+
 		highlighted, err := highlight(filename)
 		if err != nil {
 			log.Warn("Highlighting failed: ", err)

m/reader_test.go

@@ -326,9 +326,11 @@ func BenchmarkReaderDone(b *testing.B) {
 		if err != nil {
 			panic(err)
 		}
-		err = readMe._Wait()
-		if err != nil {
-			panic(err)
+
+		// Wait for the reader to finish
+		<-readMe.done
+		if readMe.err != nil {
+			panic(readMe.err)
 		}
 	}
 }