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  }