summaryrefslogtreecommitdiff
path: root/workhorse/internal/helper/writeafterreader_test.go
blob: 67cb3e6e542a972034856fc145f9344c67615bc9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
package helper

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"testing"
	"testing/iotest"
)

func TestBusyReader(t *testing.T) {
	testData := "test data"
	r := testReader(testData)
	br, _ := NewWriteAfterReader(r, &bytes.Buffer{})

	result, err := ioutil.ReadAll(br)
	if err != nil {
		t.Fatal(err)
	}

	if string(result) != testData {
		t.Fatalf("expected %q, got %q", testData, result)
	}
}

func TestFirstWriteAfterReadDone(t *testing.T) {
	writeRecorder := &bytes.Buffer{}
	br, cw := NewWriteAfterReader(&bytes.Buffer{}, writeRecorder)
	if _, err := io.Copy(ioutil.Discard, br); err != nil {
		t.Fatalf("copy from busyreader: %v", err)
	}
	testData := "test data"
	if _, err := io.Copy(cw, testReader(testData)); err != nil {
		t.Fatalf("copy test data: %v", err)
	}
	if err := cw.Flush(); err != nil {
		t.Fatalf("flush error: %v", err)
	}
	if result := writeRecorder.String(); result != testData {
		t.Fatalf("expected %q, got %q", testData, result)
	}
}

func TestWriteDelay(t *testing.T) {
	writeRecorder := &bytes.Buffer{}
	w := &complainingWriter{Writer: writeRecorder}
	br, cw := NewWriteAfterReader(&bytes.Buffer{}, w)

	testData1 := "1 test"
	if _, err := io.Copy(cw, testReader(testData1)); err != nil {
		t.Fatalf("error on first copy: %v", err)
	}

	// Unblock the coupled writer by draining the reader
	if _, err := io.Copy(ioutil.Discard, br); err != nil {
		t.Fatalf("copy from busyreader: %v", err)
	}
	// Now it is no longer an error if 'w' receives a Write()
	w.CheerUp()

	testData2 := "2 experiment"
	if _, err := io.Copy(cw, testReader(testData2)); err != nil {
		t.Fatalf("error on second copy: %v", err)
	}

	if err := cw.Flush(); err != nil {
		t.Fatalf("flush error: %v", err)
	}

	expected := testData1 + testData2
	if result := writeRecorder.String(); result != expected {
		t.Fatalf("total write: expected %q, got %q", expected, result)
	}
}

func TestComplainingWriterSanity(t *testing.T) {
	recorder := &bytes.Buffer{}
	w := &complainingWriter{Writer: recorder}

	testData := "test data"
	if _, err := io.Copy(w, testReader(testData)); err == nil {
		t.Error("error expected, none received")
	}

	w.CheerUp()
	if _, err := io.Copy(w, testReader(testData)); err != nil {
		t.Errorf("copy after CheerUp: %v", err)
	}

	if result := recorder.String(); result != testData {
		t.Errorf("expected %q, got %q", testData, result)
	}
}

func testReader(data string) io.Reader {
	return iotest.OneByteReader(bytes.NewBuffer([]byte(data)))
}

type complainingWriter struct {
	happy bool
	io.Writer
}

func (comp *complainingWriter) Write(data []byte) (int, error) {
	if comp.happy {
		return comp.Writer.Write(data)
	}

	return 0, fmt.Errorf("I am unhappy about you wanting to write %q", data)
}

func (comp *complainingWriter) CheerUp() {
	comp.happy = true
}