LogDog: Add Butler stream server package.

client/logdog/butlerlib/streamclient/client_test.go

Index: client/logdog/butlerlib/streamclient/client_test.go
diff --git a/client/logdog/butlerlib/streamclient/client_test.go b/client/logdog/butlerlib/streamclient/client_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..0faa40d1a5a508127030fae91892c3b859da255b
--- /dev/null
+++ b/client/logdog/butlerlib/streamclient/client_test.go
@@ -0,0 +1,110 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package streamclient
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/luci/luci-go/client/logdog/butlerlib/streamproto"
+	"github.com/luci/luci-go/common/clock/clockflag"
+	"github.com/luci/luci-go/common/clock/testclock"
+	"github.com/luci/luci-go/common/recordio"
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+type testStreamWriteCloser struct {
+	bytes.Buffer
+	addr string
+
+	err    error
+	closed bool
+}
+
+func (ts *testStreamWriteCloser) Write(d []byte) (int, error) {
+	if ts.err != nil {
+		return 0, ts.err
+	}
+	return ts.Buffer.Write(d)
+}
+
+func (ts *testStreamWriteCloser) Close() error {
+	ts.closed = true
+	return nil
+}
+
+func TestClient(t *testing.T) {
+	Convey(`A client registry with a test protocol`, t, func() {
+		tswcErr := error(nil)
+		tswc := (*testStreamWriteCloser)(nil)
+
+		reg := protocolRegistry{}
+		reg.register("test", func(addr string) (Client, error) {
+			return &clientImpl{
+				factory: func() (io.WriteCloser, error) {
+					tswc = &testStreamWriteCloser{
+						addr: addr,
+						err:  tswcErr,
+					}
+					return tswc, nil
+				},
+			}, nil
+		})
+
+		flags := streamproto.Flags{
+			Name:      "test",
+			Timestamp: clockflag.Time(testclock.TestTimeUTC),
+		}
+
+		Convey(`Will panic if the same protocol is registered twice.`, func() {
+			So(func() { reg.register("test", nil) }, ShouldPanic)
+			So(func() { reg.register("test2", nil) }, ShouldNotPanic)
+		})
+
+		Convey(`Will fail to instantiate a Client with an invalid protocol.`, func() {
+			_, err := reg.newClient("fake:foo")
+			So(err, ShouldNotBeNil)
+		})
+
+		Convey(`Can instantiate a new client.`, func() {
+			client, err := reg.newClient("test:foo")
+			So(err, ShouldBeNil)
+			So(client, ShouldHaveSameTypeAs, &clientImpl{})
+
+			Convey(`That can instantiate new Streams.`, func() {
+				stream, err := client.NewStream(flags)
+				So(err, ShouldBeNil)
+				So(stream, ShouldHaveSameTypeAs, &streamImpl{})
+
+				si := stream.(*streamImpl)
+				So(si.WriteCloser, ShouldHaveSameTypeAs, &testStreamWriteCloser{})
+
+				tswc := si.WriteCloser.(*testStreamWriteCloser)
+				So(tswc.addr, ShouldEqual, "foo")
+
+				Convey(`The stream should have the stream header written to it.`, func() {
+					So(tswc.Next(len(streamproto.ProtocolFrameHeaderMagic)), ShouldResemble,
+						streamproto.ProtocolFrameHeaderMagic)
+
+					r := recordio.NewReader(tswc, -1)
+					f, err := r.ReadFrameAll()
+					So(err, ShouldBeNil)
+					So(string(f), ShouldResemble, `{"name":"test","timestamp":"0001-02-03T04:05:06.000000007Z"}`)
+				})
+			})
+
+			Convey(`If the stream fails to write the handshake, it will be closed.`, func() {
+				tswcErr = errors.New("test error")
+				_, err := client.NewStream(flags)
+				So(err, ShouldNotBeNil)
+
+				So(tswc, ShouldNotBeNil)
+				So(tswc.closed, ShouldBeTrue)
+			})
+		})
+	})
+}