1package helper
2
3import (
4	"bytes"
5	"fmt"
6	"io"
7	"io/ioutil"
8	"testing"
9	"testing/iotest"
10)
11
12func TestBusyReader(t *testing.T) {
13	testData := "test data"
14	r := testReader(testData)
15	br, _ := NewWriteAfterReader(r, &bytes.Buffer{})
16
17	result, err := ioutil.ReadAll(br)
18	if err != nil {
19		t.Fatal(err)
20	}
21
22	if string(result) != testData {
23		t.Fatalf("expected %q, got %q", testData, result)
24	}
25}
26
27func TestFirstWriteAfterReadDone(t *testing.T) {
28	writeRecorder := &bytes.Buffer{}
29	br, cw := NewWriteAfterReader(&bytes.Buffer{}, writeRecorder)
30	if _, err := io.Copy(ioutil.Discard, br); err != nil {
31		t.Fatalf("copy from busyreader: %v", err)
32	}
33	testData := "test data"
34	if _, err := io.Copy(cw, testReader(testData)); err != nil {
35		t.Fatalf("copy test data: %v", err)
36	}
37	if err := cw.Flush(); err != nil {
38		t.Fatalf("flush error: %v", err)
39	}
40	if result := writeRecorder.String(); result != testData {
41		t.Fatalf("expected %q, got %q", testData, result)
42	}
43}
44
45func TestWriteDelay(t *testing.T) {
46	writeRecorder := &bytes.Buffer{}
47	w := &complainingWriter{Writer: writeRecorder}
48	br, cw := NewWriteAfterReader(&bytes.Buffer{}, w)
49
50	testData1 := "1 test"
51	if _, err := io.Copy(cw, testReader(testData1)); err != nil {
52		t.Fatalf("error on first copy: %v", err)
53	}
54
55	// Unblock the coupled writer by draining the reader
56	if _, err := io.Copy(ioutil.Discard, br); err != nil {
57		t.Fatalf("copy from busyreader: %v", err)
58	}
59	// Now it is no longer an error if 'w' receives a Write()
60	w.CheerUp()
61
62	testData2 := "2 experiment"
63	if _, err := io.Copy(cw, testReader(testData2)); err != nil {
64		t.Fatalf("error on second copy: %v", err)
65	}
66
67	if err := cw.Flush(); err != nil {
68		t.Fatalf("flush error: %v", err)
69	}
70
71	expected := testData1 + testData2
72	if result := writeRecorder.String(); result != expected {
73		t.Fatalf("total write: expected %q, got %q", expected, result)
74	}
75}
76
77func TestComplainingWriterSanity(t *testing.T) {
78	recorder := &bytes.Buffer{}
79	w := &complainingWriter{Writer: recorder}
80
81	testData := "test data"
82	if _, err := io.Copy(w, testReader(testData)); err == nil {
83		t.Error("error expected, none received")
84	}
85
86	w.CheerUp()
87	if _, err := io.Copy(w, testReader(testData)); err != nil {
88		t.Errorf("copy after CheerUp: %v", err)
89	}
90
91	if result := recorder.String(); result != testData {
92		t.Errorf("expected %q, got %q", testData, result)
93	}
94}
95
96func testReader(data string) io.Reader {
97	return iotest.OneByteReader(bytes.NewBuffer([]byte(data)))
98}
99
100type complainingWriter struct {
101	happy bool
102	io.Writer
103}
104
105func (comp *complainingWriter) Write(data []byte) (int, error) {
106	if comp.happy {
107		return comp.Writer.Write(data)
108	}
109
110	return 0, fmt.Errorf("I am unhappy about you wanting to write %q", data)
111}
112
113func (comp *complainingWriter) CheerUp() {
114	comp.happy = true
115}
116