github.com/rogpeppe/juju@v0.0.0-20140613142852-6337964b789e/environs/sshstorage/linewrapwriter_test.go (about) 1 // Copyright 2013 Canonical Ltd. 2 // Licensed under the AGPLv3, see LICENCE file for details. 3 4 package sshstorage_test 5 6 import ( 7 "bytes" 8 "errors" 9 "io" 10 11 gc "launchpad.net/gocheck" 12 13 "github.com/juju/juju/environs/sshstorage" 14 ) 15 16 type wrapWriterSuite struct{} 17 18 var _ = gc.Suite(&wrapWriterSuite{}) 19 20 func (*wrapWriterSuite) TestLineWrapWriterBadLength(c *gc.C) { 21 var buf bytes.Buffer 22 c.Assert(func() { sshstorage.NewLineWrapWriter(&buf, 0) }, gc.PanicMatches, "lineWrapWriter with line length <= 0") 23 c.Assert(func() { sshstorage.NewLineWrapWriter(&buf, -1) }, gc.PanicMatches, "lineWrapWriter with line length <= 0") 24 } 25 26 func (*wrapWriterSuite) TestLineWrapWriter(c *gc.C) { 27 type test struct { 28 input []string 29 lineLength int 30 expected string 31 } 32 tests := []test{{ 33 input: []string{""}, 34 lineLength: 1, 35 expected: "", 36 }, { 37 input: []string{"hi!"}, 38 lineLength: 1, 39 expected: "h\ni\n!\n", 40 }, { 41 input: []string{"hi!"}, 42 lineLength: 2, 43 // Note: no trailing newline. 44 expected: "hi\n!", 45 }, { 46 input: []string{"", "h", "i!"}, 47 lineLength: 2, 48 expected: "hi\n!", 49 }, { 50 input: []string{"", "h", "i!"}, 51 lineLength: 2, 52 expected: "hi\n!", 53 }, { 54 input: []string{"hi", "!!"}, 55 lineLength: 2, 56 expected: "hi\n!!\n", 57 }, { 58 input: []string{"hi", "!", "!"}, 59 lineLength: 2, 60 expected: "hi\n!!\n", 61 }, { 62 input: []string{"h", "i", "!!"}, 63 lineLength: 2, 64 expected: "hi\n!!\n", 65 }} 66 for i, t := range tests { 67 c.Logf("test %d: %q, line length %d", i, t.input, t.lineLength) 68 var buf bytes.Buffer 69 w := sshstorage.NewLineWrapWriter(&buf, t.lineLength) 70 c.Assert(w, gc.NotNil) 71 for _, input := range t.input { 72 n, err := w.Write([]byte(input)) 73 c.Assert(err, gc.IsNil) 74 c.Assert(n, gc.Equals, len(input)) 75 } 76 c.Assert(buf.String(), gc.Equals, t.expected) 77 } 78 } 79 80 type limitedWriter struct { 81 io.Writer 82 remaining int 83 } 84 85 var writeLimited = errors.New("write limited") 86 87 func (w *limitedWriter) Write(buf []byte) (int, error) { 88 inputlen := len(buf) 89 if len(buf) > w.remaining { 90 buf = buf[:w.remaining] 91 } 92 n, err := w.Writer.Write(buf) 93 w.remaining -= n 94 if n < inputlen && err == nil { 95 err = writeLimited 96 } 97 return n, err 98 } 99 100 func (*wrapWriterSuite) TestLineWrapWriterErrors(c *gc.C) { 101 // Note: after an error is returned, all bets are off. 102 // In the only place we use this code, we bail out immediately. 103 const lineLength = 3 104 type test struct { 105 input string 106 output string 107 limit int 108 written int 109 err error 110 } 111 tests := []test{{ 112 input: "abc", 113 output: "abc", 114 limit: 3, // "\n" will be limited 115 written: 3, 116 err: writeLimited, 117 }, { 118 input: "abc", 119 output: "abc\n", 120 limit: 4, 121 written: 3, // 3/3 bytes of input 122 }, { 123 input: "abc", 124 output: "ab", 125 limit: 2, 126 written: 2, // 2/3 bytes of input 127 err: writeLimited, 128 }, { 129 input: "abc!", 130 output: "abc\n", 131 limit: 4, 132 written: 3, // 3/4 bytes of input 133 err: writeLimited, 134 }} 135 for i, t := range tests { 136 c.Logf("test %d: %q, limit %d", i, t.input, t.limit) 137 var buf bytes.Buffer 138 wrapWriter := &limitedWriter{&buf, t.limit} 139 w := sshstorage.NewLineWrapWriter(wrapWriter, lineLength) 140 c.Assert(w, gc.NotNil) 141 n, err := w.Write([]byte(t.input)) 142 c.Assert(n, gc.Equals, t.written) 143 c.Assert(buf.String(), gc.Equals, t.output) 144 c.Assert(err, gc.Equals, t.err) 145 } 146 }