github.com/coreos/goproxy@v0.0.0-20190513173959-f8dc2d7ba04e/regretable/regretreader_test.go (about)

     1  package regretable_test
     2  
     3  import (
     4  	. "github.com/elazarl/goproxy/regretable"
     5  	"bytes"
     6  	"io"
     7  	"io/ioutil"
     8  	"strings"
     9  	"testing"
    10  )
    11  
    12  func TestRegretableReader(t *testing.T) {
    13  	buf := new(bytes.Buffer)
    14  	mb := NewRegretableReader(buf)
    15  	word := "12345678"
    16  	buf.WriteString(word)
    17  
    18  	fivebytes := make([]byte, 5)
    19  	mb.Read(fivebytes)
    20  	mb.Regret()
    21  
    22  	s, _ := ioutil.ReadAll(mb)
    23  	if string(s) != word {
    24  		t.Errorf("Uncommited read is gone, [%d,%d] actual '%v' expected '%v'\n", len(s), len(word), string(s), word)
    25  	}
    26  }
    27  
    28  func TestRegretableEmptyRead(t *testing.T) {
    29  	buf := new(bytes.Buffer)
    30  	mb := NewRegretableReader(buf)
    31  	word := "12345678"
    32  	buf.WriteString(word)
    33  
    34  	zero := make([]byte, 0)
    35  	mb.Read(zero)
    36  	mb.Regret()
    37  
    38  	s, err := ioutil.ReadAll(mb)
    39  	if string(s) != word {
    40  		t.Error("Uncommited read is gone, actual:", string(s), "expected:", word, "err:", err)
    41  	}
    42  }
    43  
    44  func TestRegretableAlsoEmptyRead(t *testing.T) {
    45  	buf := new(bytes.Buffer)
    46  	mb := NewRegretableReader(buf)
    47  	word := "12345678"
    48  	buf.WriteString(word)
    49  
    50  	one := make([]byte, 1)
    51  	zero := make([]byte, 0)
    52  	five := make([]byte, 5)
    53  	mb.Read(one)
    54  	mb.Read(zero)
    55  	mb.Read(five)
    56  	mb.Regret()
    57  
    58  	s, _ := ioutil.ReadAll(mb)
    59  	if string(s) != word {
    60  		t.Error("Uncommited read is gone", string(s), "expected", word)
    61  	}
    62  }
    63  
    64  func TestRegretableRegretBeforeRead(t *testing.T) {
    65  	buf := new(bytes.Buffer)
    66  	mb := NewRegretableReader(buf)
    67  	word := "12345678"
    68  	buf.WriteString(word)
    69  
    70  	five := make([]byte, 5)
    71  	mb.Regret()
    72  	mb.Read(five)
    73  
    74  	s, err := ioutil.ReadAll(mb)
    75  	if string(s) != "678" {
    76  		t.Error("Uncommited read is gone", string(s), len(string(s)), "expected", "678", len("678"), "err:", err)
    77  	}
    78  }
    79  
    80  func TestRegretableFullRead(t *testing.T) {
    81  	buf := new(bytes.Buffer)
    82  	mb := NewRegretableReader(buf)
    83  	word := "12345678"
    84  	buf.WriteString(word)
    85  
    86  	twenty := make([]byte, 20)
    87  	mb.Read(twenty)
    88  	mb.Regret()
    89  
    90  	s, _ := ioutil.ReadAll(mb)
    91  	if string(s) != word {
    92  		t.Error("Uncommited read is gone", string(s), len(string(s)), "expected", word, len(word))
    93  	}
    94  }
    95  
    96  func assertEqual(t *testing.T, expected, actual string) {
    97  	if expected!=actual {
    98  		t.Fatal("Expected", expected, "actual", actual)
    99  	}
   100  }
   101  
   102  func assertReadAll(t *testing.T, r io.Reader) string {
   103  	s, err := ioutil.ReadAll(r)
   104  	if err!=nil {
   105  		t.Fatal("error when reading", err)
   106  	}
   107  	return string(s)
   108  }
   109  
   110  func TestRegretableRegretTwice(t *testing.T) {
   111  	buf := new(bytes.Buffer)
   112  	mb := NewRegretableReader(buf)
   113  	word := "12345678"
   114  	buf.WriteString(word)
   115  
   116  	assertEqual(t, word, assertReadAll(t, mb))
   117  	mb.Regret()
   118  	assertEqual(t, word, assertReadAll(t, mb))
   119  	mb.Regret()
   120  	assertEqual(t, word, assertReadAll(t, mb))
   121  }
   122  
   123  type CloseCounter struct {
   124  	r      io.Reader
   125  	closed int
   126  }
   127  
   128  func (cc *CloseCounter) Read(b []byte) (int, error) {
   129  	return cc.r.Read(b)
   130  }
   131  
   132  func (cc *CloseCounter) Close() error {
   133  	cc.closed++
   134  	return nil
   135  }
   136  
   137  func assert(t *testing.T, b bool, msg string) {
   138  	if !b {
   139  		t.Errorf("Assertion Error: %s", msg)
   140  	}
   141  }
   142  
   143  func TestRegretableCloserSizeRegrets(t *testing.T) {
   144  	defer func() {
   145  		if r := recover(); r == nil || !strings.Contains(r.(string), "regret") {
   146  			t.Error("Did not panic when regretting overread buffer:", r)
   147  		}
   148  	}()
   149  	buf := new(bytes.Buffer)
   150  	buf.WriteString("123456")
   151  	mb := NewRegretableReaderCloserSize(ioutil.NopCloser(buf), 3)
   152  	mb.Read(make([]byte, 4))
   153  	mb.Regret()
   154  }
   155  
   156  func TestRegretableCloserRegretsClose(t *testing.T) {
   157  	buf := new(bytes.Buffer)
   158  	cc := &CloseCounter{buf, 0}
   159  	mb := NewRegretableReaderCloser(cc)
   160  	word := "12345678"
   161  	buf.WriteString(word)
   162  
   163  	mb.Read([]byte{0})
   164  	mb.Close()
   165  	if cc.closed != 1 {
   166  		t.Error("RegretableReaderCloser ignores Close")
   167  	}
   168  	mb.Regret()
   169  	mb.Close()
   170  	if cc.closed != 2 {
   171  		t.Error("RegretableReaderCloser does ignore Close after regret")
   172  	}
   173  	// TODO(elazar): return an error if client issues Close more than once after regret
   174  }