ioutil: Context and OnError helpers

1 files changed, 101 insertions, 1 deletions

diff --git a/utils/ioutil/common_test.go b/utils/ioutil/common_test.go
index 2d6ef80..27bfa62 100644
--- a/utils/ioutil/common_test.go
+++ b/utils/ioutil/common_test.go
@@ -2,6 +2,7 @@ package ioutil
 
 import (
 	"bytes"
+	"context"
 	"io/ioutil"
 	"strings"
 	"testing"
@@ -55,6 +56,106 @@ func (s *CommonSuite) TestNewReadCloser(c *C) {
 	c.Assert(closer.called, Equals, 1)
 }
 
+func (s *CommonSuite) TestNewContextReader(c *C) {
+	buf := bytes.NewBuffer([]byte("12"))
+	ctx, close := context.WithCancel(context.Background())
+
+	r := NewContextReader(ctx, buf)
+
+	b := make([]byte, 1)
+	n, err := r.Read(b)
+	c.Assert(n, Equals, 1)
+	c.Assert(err, IsNil)
+
+	close()
+	n, err = r.Read(b)
+	c.Assert(n, Equals, 0)
+	c.Assert(err, NotNil)
+}
+
+func (s *CommonSuite) TestNewContextReadCloser(c *C) {
+	buf := NewReadCloser(bytes.NewBuffer([]byte("12")), &closer{})
+	ctx, close := context.WithCancel(context.Background())
+
+	r := NewContextReadCloser(ctx, buf)
+
+	b := make([]byte, 1)
+	n, err := r.Read(b)
+	c.Assert(n, Equals, 1)
+	c.Assert(err, IsNil)
+
+	close()
+	n, err = r.Read(b)
+	c.Assert(n, Equals, 0)
+	c.Assert(err, NotNil)
+
+	c.Assert(r.Close(), IsNil)
+}
+
+func (s *CommonSuite) TestNewContextWriter(c *C) {
+	buf := bytes.NewBuffer(nil)
+	ctx, close := context.WithCancel(context.Background())
+
+	r := NewContextWriter(ctx, buf)
+
+	n, err := r.Write([]byte("1"))
+	c.Assert(n, Equals, 1)
+	c.Assert(err, IsNil)
+
+	close()
+	n, err = r.Write([]byte("1"))
+	c.Assert(n, Equals, 0)
+	c.Assert(err, NotNil)
+}
+
+func (s *CommonSuite) TestNewContextWriteCloser(c *C) {
+	buf := NewWriteCloser(bytes.NewBuffer(nil), &closer{})
+	ctx, close := context.WithCancel(context.Background())
+
+	w := NewContextWriteCloser(ctx, buf)
+
+	n, err := w.Write([]byte("1"))
+	c.Assert(n, Equals, 1)
+	c.Assert(err, IsNil)
+
+	close()
+	n, err = w.Write([]byte("1"))
+	c.Assert(n, Equals, 0)
+	c.Assert(err, NotNil)
+
+	c.Assert(w.Close(), IsNil)
+}
+
+func (s *CommonSuite) TestNewWriteCloserOnError(c *C) {
+	buf := NewWriteCloser(bytes.NewBuffer(nil), &closer{})
+
+	ctx, close := context.WithCancel(context.Background())
+
+	var called error
+	w := NewWriteCloserOnError(NewContextWriteCloser(ctx, buf), func(err error) {
+		called = err
+	})
+
+	close()
+	w.Write(nil)
+
+	c.Assert(called, NotNil)
+}
+
+func (s *CommonSuite) TestNewReadCloserOnError(c *C) {
+	buf := NewReadCloser(bytes.NewBuffer(nil), &closer{})
+	ctx, close := context.WithCancel(context.Background())
+
+	var called error
+	w := NewReadCloserOnError(NewContextReadCloser(ctx, buf), func(err error) {
+		called = err
+	})
+
+	close()
+	w.Read(nil)
+
+	c.Assert(called, NotNil)
+}
 func ExampleCheckClose() {
 	// CheckClose is commonly used with named return values
 	f := func() (err error) {
@@ -68,7 +169,6 @@ func ExampleCheckClose() {
 
 		// if err is not nil, CheckClose will assign any close errors to it
 		return err
-
 	}
 
 	err := f()