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