github.com/u-root/u-root@v7.0.1-0.20200915234505-ad7babab0a8e+incompatible/pkg/ts/ts_test.go (about)

     1  // Copyright 2020 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package ts
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"io"
    11  	"io/ioutil"
    12  	"strings"
    13  	"testing"
    14  	"time"
    15  )
    16  
    17  func TestPrependTimestamp(t *testing.T) {
    18  	format := func(time.Time) string { return "#" }
    19  	tests := []struct {
    20  		name, input, want string
    21  	}{
    22  		{
    23  			name:  "empty",
    24  			input: "",
    25  			want:  "",
    26  		},
    27  		{
    28  			name:  "single blank line",
    29  			input: "\n",
    30  			want:  "#\n",
    31  		},
    32  		{
    33  			name:  "blank lines",
    34  			input: "\n\n\n\n",
    35  			want:  "#\n#\n#\n#\n",
    36  		},
    37  		{
    38  			name:  "text",
    39  			input: "hello\nworld\n\n\n",
    40  			want:  "#hello\n#world\n#\n#\n",
    41  		},
    42  	}
    43  
    44  	for _, tt := range tests {
    45  		t.Run(tt.name, func(t *testing.T) {
    46  			pt := &PrependTimestamp{
    47  				R:      bytes.NewBufferString(tt.input),
    48  				Format: format,
    49  			}
    50  			got := &bytes.Buffer{}
    51  			if _, err := io.Copy(got, pt); err != nil {
    52  				t.Errorf("io.Copy returned an error: %v", err)
    53  			}
    54  			if !bytes.Equal(got.Bytes(), []byte(tt.want)) {
    55  				t.Errorf("PrependTimestamp = %q; want %q", got.String(), tt.want)
    56  			}
    57  		})
    58  	}
    59  }
    60  
    61  // TestPrependTimestampBuffering ensures two important properties with regards
    62  // to buffering which would be easy to miss with just a readline implementation:
    63  //     1. Data is printed before the whole line is available.
    64  //     2. Timestamp is generated at the beginning of the line, not at the end
    65  //        of the previous line.
    66  func TestPrependTimestampBuffering(t *testing.T) {
    67  	// Mock out the format function.
    68  	i := 0
    69  	format := func(time.Time) string {
    70  		return fmt.Sprint(i)
    71  	}
    72  
    73  	// Control exactly how many bytes are returned with each call to data.Read.
    74  	data := &bytes.Buffer{}
    75  	pt := &PrependTimestamp{
    76  		R:      data,
    77  		Format: format,
    78  	}
    79  
    80  	// These tests must run sequentially, so no tt.Run().
    81  	for _, tt := range []struct {
    82  		i        int
    83  		text     string
    84  		buffer   []byte
    85  		wantText string
    86  		wantErr  error
    87  	}{
    88  		{
    89  			// Data is printed before the whole line is available.
    90  			i:        1,
    91  			text:     "Waiting...",
    92  			buffer:   make([]byte, 100),
    93  			wantErr:  nil,
    94  			wantText: "1Waiting...",
    95  		},
    96  		{
    97  			text:     "DONE\n",
    98  			buffer:   make([]byte, 100),
    99  			wantErr:  nil,
   100  			wantText: "DONE\n",
   101  		},
   102  		{
   103  			// Timestamp is generated at the beginning of the line.
   104  			i:        2,
   105  			text:     "Hello",
   106  			buffer:   make([]byte, 2),
   107  			wantErr:  nil,
   108  			wantText: "2H",
   109  		},
   110  		{
   111  			text:     "",
   112  			buffer:   make([]byte, 2),
   113  			wantErr:  nil,
   114  			wantText: "el",
   115  		},
   116  		{
   117  			text:     "",
   118  			buffer:   make([]byte, 2),
   119  			wantErr:  nil,
   120  			wantText: "lo",
   121  		},
   122  		{
   123  			text:     "",
   124  			buffer:   make([]byte, 2),
   125  			wantErr:  io.EOF,
   126  			wantText: "",
   127  		},
   128  	} {
   129  		data.Write([]byte(tt.text))
   130  		i = tt.i
   131  		n, err := pt.Read(tt.buffer)
   132  		if err != tt.wantErr {
   133  			t.Errorf("PrependTimestamp.Read(%q) err = %v; want %v",
   134  				tt.text, err, tt.wantErr)
   135  		}
   136  		if !bytes.Equal(tt.buffer[:n], []byte(tt.wantText)) {
   137  			t.Errorf("PrependTimestamp.Read(%q) buffer = %q; want %q",
   138  				tt.text, tt.buffer[:n], tt.wantText)
   139  		}
   140  	}
   141  }
   142  
   143  // BenchmarkPrependTime measures the throughput of PrependTime where N is
   144  // measured in bytes.
   145  func BenchmarkPrependTime(b *testing.B) {
   146  	line := "hello world\n"
   147  	data := strings.Repeat(line, (b.N+len(line))/len(line))[:b.N]
   148  	pt := New(bytes.NewBufferString(data))
   149  	b.ResetTimer()
   150  	if _, err := io.Copy(ioutil.Discard, pt); err != nil {
   151  		b.Fatal(err)
   152  	}
   153  }