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