Repository: ably/ably-go
Pull request #185 (2020-08-16 10:37:54)

RTL13: Handle DETACHED while not DETACHING.
[RTL13](https://docs.ably.io/client-lib-development-guide/features/#RTL13)

Missing SUSPENDED, which isn't implemented.


Patch:
diff --git a/ably/internal/ablyutil/time.go b/ably/internal/ablyutil/time.go
new file mode 100644
index 0000000..dfd6d0c
--- /dev/null
+++ b/ably/internal/ablyutil/time.go
@@ -0,0 +1,27 @@
+package ablyutil
+
+import (
+	"context"
+	"time"
+)
+
+// After returns a channel that is sent the current time once the given
+// duration has passed. If the context is cancelled before that, the channel
+// is immediately closed.
+func After(ctx context.Context, d time.Duration) <-chan time.Time {
+	timer := time.NewTimer(d)
+
+	ch := make(chan time.Time)
+
+	go func() {
+		defer timer.Stop()
+		select {
+		case <-ctx.Done():
+			close(ch)
+		case t := <-timer.C:
+			ch <- t
+		}
+	}()
+
+	return ch
+}
diff --git a/ably/options.go b/ably/options.go
index e41cf90..5ce7ccc 100644
--- a/ably/options.go
+++ b/ably/options.go
@@ -30,11 +30,13 @@ var defaultOptions = clientOptions{
 	TimeoutDisconnect:        30 * time.Second,
 	RealtimeRequestTimeout:   10 * time.Second, // DF1b
 	DisconnectedRetryTimeout: 15 * time.Second, // TO3l1
+	ChannelRetryTimeout:      15 * time.Second, // TO3l7
 	FallbackRetryTimeout:     10 * time.Minute,
 	IdempotentRestPublishing: false,
 	Port:                     80,
 	TLSPort:                  443,
 	Now:                      time.Now,
+	After:                    ablyutil.After,
 }
 
 func DefaultFallbackHosts() []string {
@@ -202,6 +204,7 @@ type clientOptions struct {
 	// DisconnectedRetryTimeout is the time to wait after a disconnection before
 	// attempting an automatic reconnection, if still disconnected.
 	DisconnectedRetryTimeout time.Duration
+	ChannelRetryTimeout      time.Duration
 
 	// Dial specifies the dial function for creating message connections used
 	// by Realtime.
@@ -223,7 +226,8 @@ type clientOptions struct {
 	Trace *httptrace.ClientTrace
 
 	// Now returns the time the library should take as current.
-	Now func() time.Time
+	Now   func() time.Time
+	After func(context.Context, time.Duration) <-chan time.Time
 }
 
 func (opts *clientOptions) timeoutConnect() time.Duration {
@@ -619,6 +623,12 @@ func (os ClientOptions) DisconnectedRetryTimeout(d time.Duration) ClientOptions
 	})
 }
 
+func (os ClientOptions) ChannelRetryTimeout(d time.Duration) ClientOptions {
+	return append(os, func(os *clientOptions) {
+		os.ChannelRetryTimeout = d
+	})
+}
+
 func (os ClientOptions) HTTPMaxRetryCount(count int) ClientOptions {
 	return append(os, func(os *clientOptions) {
 		os.HTTPMaxRetryCount = count
diff --git a/ably/realtime_channel.go b/ably/realtime_channel.go
index 4b49dd4..de362d5 100644
--- a/ably/realtime_channel.go
+++ b/ably/realtime_channel.go
@@ -1,6 +1,7 @@
 package ably
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"sort"
@@ -198,10 +199,14 @@ func (c *RealtimeChannel) mayAttach(result, checkActive bool) (Result, error) {
 			return nopResult, nil
 		}
 	}
+	return c.lockAttach(result, nil)
+}
+
+func (c *RealtimeChannel) lockAttach(result bool, err error) (Result, error) {
 	if !c.client.Connection.lockIsActive() {
 		return nil, c.state.set(StateChanFailed, errAttach)
 	}
-	c.state.set(StateChanAttaching, nil)
+	c.state.set(StateChanAttaching, err)
 	var res Result
 	if result {
 		res = c.state.listenResult(attachResultStates...)
@@ -210,7 +215,7 @@ func (c *RealtimeChannel) mayAttach(result, checkActive bool) (Result, error) {
 		Action:  proto.ActionAttach,
 		Channel: c.state.channel,
 	}
-	err := c.client.Connection.send(msg, nil)
+	err = c.client.Connection.send(msg, nil)
 	if err != nil {
 		return nil, c.state.set(StateChanFailed, err)
 	}
@@ -397,7 +402,44 @@ func (c *RealtimeChannel) notify(msg *proto.ProtocolMessage) {
 		c.state.syncSet(StateChanAttached, nil)
 		c.queue.Flush()
 	case proto.ActionDetached:
-		c.state.syncSet(StateChanDetached, nil)
+		c.state.Lock()
+
+		err := error(newErrorProto(msg.Error))
+		switch c.state.current {
+		case StateChanDetaching:
+			c.state.set(StateChanDetached, err)
+			c.state.Unlock()
+			return
+		case StateChanAttached: // TODO: Also SUSPENDED; RTL13a
+			var res Result
+			res, err = c.lockAttach(true, err)
+			if err != nil {
+				break
+			}
+
+			c.state.Unlock()
+			go func() {
+				// We need to wait in another goroutine to allow more messages
+				// to reach the connection.
+
+				err = res.Wait()
+				if err == nil {
+					return
+				}
+				c.state.Lock()
+
+				c.lockStartRetryAttachLoop(err)
+			}()
+			return
+		case
+			StateChanAttaching,
+			StateChanDetached: // TODO: Should be SUSPENDED
+		default:
+			c.state.Unlock()
+			return
+		}
+
+		c.lockStartRetryAttachLoop(err)
 	case proto.ActionSync:
 		c.Presence.processIncomingMessage(msg, syncSerial(msg))
 	case proto.ActionPresence:
@@ -411,6 +453,48 @@ func (c *RealtimeChannel) notify(msg *proto.ProtocolMessage) {
 	}
 }
 
+func (c *RealtimeChannel) lockStartRetryAttachLoop(err error) {
+	// TODO: Move to SUSPENDED; move it to DETACHED for now.
+	c.state.set(StateChanDetached, err)
+	c.state.Unlock()
+
+	go func() {
+		// TODO: The listener should be set before unlocking the state.
+		// Will do once we have an EventEmitter.
+		stateChange := make(chan State, 10)
+		c.state.on(stateChange)
+		defer c.state.off(stateChange)
+
+		for !c.retryAttach(stateChange) {
+		}
+	}()
+}
+
+func (c *RealtimeChannel) retryAttach(stateChange chan State) (done bool) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	select {
+	case <-c.opts().After(ctx, c.opts().ChannelRetryTimeout):
+	case <-stateChange:
+		// Any concurrent state change cancels the retry.
+		return true
+	}
+
+	if c.client.Connection.State() != ConnectionStateConnected {
+		// RTL13c: If no longer CONNECTED, RTL3 takes over.
+		return true
+	}
+
+	err := wait(c.mayAttach(true, false))
+	if err == nil {
+		return true
+	}
+	// TODO: Move to SUSPENDED; move it to DETACHED for now.
+	c.state.syncSet(StateChanDetached, err)
+	return false
+}
+
 func (c *RealtimeChannel) isActive() bool {
 	return c.state.current == StateChanAttaching || c.state.current == StateChanAttached
 }
diff --git a/ably/realtime_conn.go b/ably/realtime_conn.go
index 8da7c9a..db4800a 100644
--- a/ably/realtime_conn.go
+++ b/ably/realtime_conn.go
@@ -326,10 +326,10 @@ func (c *Connection) Serial() int64 {
 }
 
 // State returns current state of the connection.
-func (c *Connection) State() StateEnum {
+func (c *Connection) State() ConnectionState {
 	c.state.Lock()
 	defer c.state.Unlock()
-	return c.state.current
+	return mapOldToNewConnState(c.state.current)
 }
 
 // onState relays request connection states to the given channel; onState state transition
diff --git a/ably/state.go b/ably/state.go
index 7a0ed77..88bf910 100644
--- a/ably/state.go
+++ b/ably/state.go
@@ -219,7 +219,7 @@ func (s *stateEmitter) set(state StateEnum, err error) error {
 	changed := s.current != state
 	s.current = state
 	s.err = stateError(state, err)
-	if changed {
+	if changed || err != nil {
 		s.emit(State{
 			Channel: s.channel,
 			Err:     s.err,


Test patch:
diff --git a/ably/ablytest/recorders.go b/ably/ablytest/recorders.go
index e6c11c7..bd56dbb 100644
--- a/ably/ablytest/recorders.go
+++ b/ably/ablytest/recorders.go
@@ -332,7 +332,7 @@ func (pc pipeConn) Receive(deadline time.Time) (*proto.ProtocolMessage, error) {
 	}
 	select {
 	case m, ok := <-pc.in:
-		if !ok {
+		if !ok || m == nil {
 			return nil, io.EOF
 		}
 		return m, nil
@@ -602,9 +602,9 @@ type realtimeIOCloser struct {
 func (c realtimeIOCloser) Close() error {
 	switch c.c.Connection.State() {
 	case
-		ably.StateConnInitialized,
-		ably.StateConnClosed,
-		ably.StateConnFailed:
+		ably.ConnectionStateInitialized,
+		ably.ConnectionStateClosed,
+		ably.ConnectionStateFailed:
 
 		return c.c.Connection.ErrorReason()
 	}
diff --git a/ably/auth_test.go b/ably/auth_test.go
index b8f231a..f995386 100644
--- a/ably/auth_test.go
+++ b/ably/auth_test.go
@@ -696,8 +696,8 @@ func TestAuth_ClientID(t *testing.T) {
 	if err = checkError(40012, err); err != nil {
 		t.Fatal(err)
 	}
-	if state := client.Connection.State(); state != ably.StateConnFailed {
-		t.Fatalf("want state=%q; got %q", ably.StateConnFailed, state)
+	if state := client.Connection.State(); state != ably.ConnectionStateFailed {
+		t.Fatalf("want state=%q; got %q", ably.ConnectionStateFailed, state)
 	}
 	select {
 	case state := <-failed:
diff --git a/ably/export_test.go b/ably/export_test.go
index 24c2e70..69eb08b 100644
--- a/ably/export_test.go
+++ b/ably/export_test.go
@@ -1,6 +1,7 @@
 package ably
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptrace"
 	"time"
@@ -95,6 +96,12 @@ func (c *Connection) OffState(ch chan<- State, states ...StateEnum) {
 	c.offState(ch, states...)
 }
 
+func (c *Connection) StateEnum() StateEnum {
+	c.state.Lock()
+	defer c.state.Unlock()
+	return c.state.current
+}
+
 func (c *Connection) RemoveKey() {
 	c.state.Lock()
 	defer c.state.Unlock()
@@ -131,6 +138,12 @@ func (os ClientOptions) Now(now func() time.Time) ClientOptions {
 	})
 }
 
+func (os ClientOptions) After(after func(context.Context, time.Duration) <-chan time.Time) ClientOptions {
+	return append(os, func(os *clientOptions) {
+		os.After = after
+	})
+}
+
 func (os ClientOptions) ApplyWithDefaults() *clientOptions {
 	return os.applyWithDefaults()
 }
diff --git a/ably/internal/ablyutil/time_test.go b/ably/internal/ablyutil/time_test.go
new file mode 100644
index 0000000..9478f1d
--- /dev/null
+++ b/ably/internal/ablyutil/time_test.go
@@ -0,0 +1,66 @@
+package ablyutil_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/ably/ably-go/ably/internal/ablyutil"
+	"golang.org/x/net/context"
+)
+
+func TestAfterOK(t *testing.T) {
+	t.Parallel()
+
+	const wait = 10 * time.Millisecond
+
+	ctx, cancel := context.WithTimeout(context.Background(), wait*2)
+	defer cancel()
+
+	now := time.Now()
+
+	select {
+	case firedAt := <-ablyutil.After(ctx, wait):
+		if elapsed := firedAt.Sub(now); !isCloseTo(elapsed, wait) {
+			t.Errorf("expected timer to fire after ~%v; got %v", wait, elapsed)
+		}
+	case <-ctx.Done():
+		t.Error("expected timer to be done before the context is canceled")
+	}
+}
+
+func TestAfterCanceled(t *testing.T) {
+	t.Parallel()
+
+	const wait = 10 * time.Millisecond
+
+	testCtx, cancel := context.WithTimeout(context.Background(), wait*2)
+	defer cancel()
+
+	ctx, cancel := context.WithCancel(testCtx)
+	defer cancel()
+
+	var canceledAt time.Time
+
+	go func() {
+		time.Sleep(wait / 2)
+		canceledAt = time.Now()
+		cancel()
+	}()
+
+	select {
+	case _, ok := <-ablyutil.After(ctx, wait):
+		if ok {
+			t.Error("expected timer channel to be closed on cancel")
+		}
+		if sinceCancel := time.Since(canceledAt); !isCloseTo(sinceCancel, 0) {
+			t.Errorf("expected timer to fire immediately after cancel; got %v", sinceCancel)
+		}
+	case <-testCtx.Done():
+		t.Error("expected timer to be done before the context is canceled")
+	}
+}
+
+func isCloseTo(d time.Duration, to time.Duration) bool {
+	const leeway = 1 * time.Millisecond
+	return d > to-leeway && d < to+leeway
+}
diff --git a/ably/realtime_channel_spec_test.go b/ably/realtime_channel_spec_test.go
new file mode 100644
index 0000000..41967da
--- /dev/null
+++ b/ably/realtime_channel_spec_test.go
@@ -0,0 +1,265 @@
+package ably_test
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/ably/ably-go/ably"
+	"github.com/ably/ably-go/ably/ablytest"
+	"github.com/ably/ably-go/ably/proto"
+)
+
+func TestRealtimeChannel_RTL13_HandleDetached(t *testing.T) {
+	t.Parallel()
+
+	const channelRetryTimeout = 123 * time.Millisecond
+
+	type afterCall struct {
+		d   time.Duration
+		ret chan<- time.Time
+	}
+
+	setup := func(t *testing.T) (
+		in, out chan *proto.ProtocolMessage,
+		c *ably.Realtime,
+		channel *ably.RealtimeChannel,
+		stateChanges chan ably.State,
+		afterCalls chan afterCall,
+	) {
+		in = make(chan *proto.ProtocolMessage, 1)
+		out = make(chan *proto.ProtocolMessage, 16)
+		afterCalls = make(chan afterCall, 1)
+
+		c, _ = ably.NewRealtime(ably.ClientOptions{}.
+			Token("fake:token").
+			AutoConnect(false).
+			After(func(ctx context.Context, d time.Duration) <-chan time.Time {
+				ch := make(chan time.Time)
+				afterCalls <- afterCall{d: d, ret: ch}
+				return ch
+			}).
+			ChannelRetryTimeout(channelRetryTimeout).
+			Dial(ablytest.MessagePipe(in, out)))
+
+		in <- &proto.ProtocolMessage{
+			Action:            proto.ActionConnected,
+			ConnectionID:      "connection-id",
+			ConnectionDetails: &proto.ConnectionDetails{},
+		}
+
+		err := ablytest.ConnWaiter(c, c.Connect, ably.ConnectionEventConnected).Wait()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		channel = c.Channels.Get("test")
+
+		_, err = channel.Attach()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		ablytest.Instantly.Recv(t, nil, out, t.Fatalf) // Consume ATTACH
+
+		stateChanges = make(chan ably.State, 10)
+		channel.On(stateChanges)
+
+		return
+	}
+
+	t.Run("RTL13a: when ATTACHED, successful reattach", func(t *testing.T) {
+		t.Parallel()
+
+		in, out, _, channel, stateChanges, afterCalls := setup(t)
+
+		// Get the channel to ATTACHED.
+
+		in <- &proto.ProtocolMessage{
+			Action:  proto.ActionAttached,
+			Channel: channel.Name,
+		}
+
+		var change ably.State
+
+		ablytest.Instantly.Recv(t, &change, stateChanges, t.Fatalf)
+		if expected, got := ably.StateChanAttached, change.State; expected != got {
+			t.Fatalf("expected %v; got %v (event: %+v)", expected, got, change)
+		}
+
+		errInfo := proto.ErrorInfo{
+			StatusCode: 500,
+			Code:       50500,
+			Message:    "fake error",
+		}
+
+		in <- &proto.ProtocolMessage{
+			Action:  proto.ActionDetached,
+			Channel: channel.Name,
+			Error:   &errInfo,
+		}
+
+		// Expect a transition to ATTACHING with the error.
+
+		ablytest.Instantly.Recv(t, &change, stateChanges, t.Fatalf)
+		if expected, got := ably.StateChanAttaching, change.State; expected != got {
+			t.Fatalf("expected %v; got %v (event: %+v)", expected, got, change)
+		}
+		if got := fmt.Sprint(change.Err); !strings.Contains(got, errInfo.Message) {
+			t.Fatalf("expected %+v; got %v (error: %+v)", errInfo, got, change.Err)
+		}
+
+		var msg *proto.ProtocolMessage
+		ablytest.Instantly.Recv(t, &msg, out, t.Fatalf)
+		if expected, got := proto.ActionAttach, msg.Action; expected != got {
+			t.Fatalf("expected %v; got %v (message: %+v)", expected, got, msg)
+		}
+
+		// TODO: Test attach failure too, which requires RTL4f.
+
+		in <- &proto.ProtocolMessage{
+			Action:  proto.ActionAttached,
+			Channel: channel.Name,
+		}
+		ablytest.Instantly.Recv(t, &change, stateChanges, t.Fatalf)
+		if expected, got := ably.StateChanAttached, change.State; expected != got {
+			t.Fatalf("expected %v; got %v (event: %+v)", expected, got, change)
+		}
+		if change.Err != nil {
+			t.Fatal(change.Err)
+		}
+
+		// Expect the retry loop to be finished.
+		ablytest.Instantly.NoRecv(t, nil, afterCalls, t.Fatalf)
+		ablytest.Instantly.NoRecv(t, nil, stateChanges, t.Fatalf)
+	})
+
+	t.Run("RTL13b: when ATTACHING", func(t *testing.T) {
+		t.Parallel()
+
+		in, out, _, channel, stateChanges, afterCalls := setup(t)
+
+		errInfo := proto.ErrorInfo{
+			StatusCode: 500,
+			Code:       50500,
+			Message:    "fake error",
+		}
+
+		in <- &proto.ProtocolMessage{
+			Action:  proto.ActionDetached,
+			Channel: channel.Name,
+			Error:   &errInfo,
+		}
+
+		// Expect a state change with the error.
+
+		var change ably.State
+		ablytest.Instantly.Recv(t, &change, stateChanges, t.Fatalf)
+		if expected, got := ably.StateChanDetached, change.State; expected != got {
+			t.Fatalf("expected %v; got %v (event: %+v)", expected, got, change)
+		}
+		if got := fmt.Sprint(change.Err); !strings.Contains(got, errInfo.Message) {
+			t.Fatalf("expected %+v; got %v (error: %+v)", errInfo, got, change.Err)
+		}
+
+		// Expect an attempt to attach after channelRetryTimeout.
+
+		var call afterCall
+		ablytest.Instantly.Recv(t, &call, afterCalls, t.Fatalf)
+		if expected, got := channelRetryTimeout, call.d; expected != got {
+			t.Fatalf("expected %v; got %v", expected, got)
+		}
+		call.ret <- time.Time{}
+
+		// Expect a transition to ATTACHING, and an ATTACH message.
+
+		ablytest.Instantly.Recv(t, &change, stateChanges, t.Fatalf)
+		if expected, got := ably.StateChanAttaching, change.State; expected != got {
+			t.Fatalf("expected %v; got %v (event: %+v)", expected, got, change)
+		}
+
+		var msg *proto.ProtocolMessage
+		ablytest.Instantly.Recv(t, &msg, out, t.Fatalf)
+		if expected, got := proto.ActionAttach, msg.Action; expected != got {
+			t.Fatalf("expected %v; got %v (message: %+v)", expected, got, msg)
+		}
+
+		// TODO: Test attach failure too, which requires RTL4f.
+
+		in <- &proto.ProtocolMessage{
+			Action:  proto.ActionAttached,
+			Channel: channel.Name,
+		}
+		ablytest.Instantly.Recv(t, &change, stateChanges, t.Fatalf)
+		if expected, got := ably.StateChanAttached, change.State; expected != got {
+			t.Fatalf("expected %v; got %v (event: %+v)", expected, got, change)
+		}
+		if change.Err != nil {
+			t.Fatal(change.Err)
+		}
+
+		// Expect the retry loop to be finished.
+		ablytest.Instantly.NoRecv(t, nil, afterCalls, t.Fatalf)
+		ablytest.Instantly.NoRecv(t, nil, stateChanges, t.Fatalf)
+	})
+
+	t.Run("RTL13c: stop on non-CONNECTED", func(t *testing.T) {
+		t.Parallel()
+
+		in, out, c, channel, stateChanges, afterCalls := setup(t)
+
+		errInfo := proto.ErrorInfo{
+			StatusCode: 500,
+			Code:       50500,
+			Message:    "fake error",
+		}
+
+		in <- &proto.ProtocolMessage{
+			Action:  proto.ActionDetached,
+			Channel: channel.Name,
+			Error:   &errInfo,
+		}
+
+		// Expect a state change with the error.
+
+		var change ably.State
+		ablytest.Instantly.Recv(t, &change, stateChanges, t.Fatalf)
+		if expected, got := ably.StateChanDetached, change.State; expected != got {
+			t.Fatalf("expected %v; got %v (event: %+v)", expected, got, change)
+		}
+		if got := fmt.Sprint(change.Err); !strings.Contains(got, errInfo.Message) {
+			t.Fatalf("expected %+v; got %v (error: %+v)", errInfo, got, change.Err)
+		}
+
+		// Expect an attempt to attach after channelRetryTimeout.
+
+		var call afterCall
+		ablytest.Instantly.Recv(t, &call, afterCalls, t.Fatalf)
+		if expected, got := channelRetryTimeout, call.d; expected != got {
+			t.Fatalf("expected %v; got %v", expected, got)
+		}
+
+		// Get the connection to a non-CONNECTED state by closing in.
+
+		err := ablytest.ConnWaiter(c, func() {
+			close(in)
+		}, ably.ConnectionEventDisconnected).Wait()
+		if !errors.Is(err, io.EOF) {
+			t.Fatal(err)
+		}
+
+		// Now trigger the channelRetryTimeout.
+
+		call.ret <- time.Time{}
+
+		// Since the connection isn't CONNECTED, the retry loop should finish.
+
+		ablytest.Instantly.NoRecv(t, nil, afterCalls, t.Fatalf)
+		ablytest.Instantly.NoRecv(t, nil, stateChanges, t.Fatalf)
+		ablytest.Instantly.NoRecv(t, nil, out, t.Fatalf)
+	})
+}
diff --git a/ably/realtime_client_test.go b/ably/realtime_client_test.go
index 814b0bd..32b5d1d 100644
--- a/ably/realtime_client_test.go
+++ b/ably/realtime_client_test.go
@@ -33,16 +33,16 @@ func TestRealtime_RealtimeHost(t *testing.T) {
 			t.Errorf("NewRealtime=%s (host=%s)", err, host)
 			continue
 		}
-		if state := client.Connection.State(); state != ably.StateConnInitialized {
-			t.Errorf("want state=%v; got %s", ably.StateConnInitialized, state)
+		if state := client.Connection.State(); state != ably.ConnectionStateInitialized {
+			t.Errorf("want state=%v; got %s", ably.ConnectionStateInitialized, state)
 			continue
 		}
 		if err := checkError(80000, ablytest.ConnWaiter(client, client.Connect, ably.ConnectionEventConnected).Wait()); err != nil {
 			t.Errorf("%s (host=%s)", err, host)
 			continue
 		}
-		if state := client.Connection.State(); state != ably.StateConnFailed {
-			t.Errorf("want state=%v; got %s", ably.StateConnFailed, state)
+		if state := client.Connection.State(); state != ably.ConnectionStateFailed {
+			t.Errorf("want state=%v; got %s", ably.ConnectionStateFailed, state)
 			continue
 		}
 		if err := checkError(80000, ablytest.FullRealtimeCloser(client).Close()); err != nil {
diff --git a/ably/realtime_conn_spec_test.go b/ably/realtime_conn_spec_test.go
index 1523043..6d16b80 100644
--- a/ably/realtime_conn_spec_test.go
+++ b/ably/realtime_conn_spec_test.go
@@ -968,7 +968,7 @@ func TestRealtimeConn_RTN15c4(t *testing.T) {
 		t.Errorf("expected %d got %d", errInfo.StatusCode, reason.StatusCode)
 	}
 	// The client should transition to the FAILED state
-	if expected, got := ably.StateConnFailed, client.Connection.State(); expected != got {
+	if expected, got := ably.ConnectionStateFailed, client.Connection.State(); expected != got {
 		t.Fatalf("expected transition to %v, got %v", expected, got)
 	}
 }
@@ -1621,7 +1621,9 @@ func TestRealtimeConn_RTN23(t *testing.T) {
 				ConnectionDetails: &connDetails,
 			}
 			dials <- u
-			return ablytest.MessagePipe(in, nil)(p, u)
+			return ablytest.MessagePipe(in, nil,
+				ablytest.MessagePipeWithNowFunc(time.Now),
+			)(p, u)
 		}))
 	err := ablytest.ConnWaiter(c, c.Connect, ably.ConnectionEventConnected).Wait()
 	if err != nil {
diff --git a/ably/realtime_conn_test.go b/ably/realtime_conn_test.go
index c13252f..08cad4e 100644
--- a/ably/realtime_conn_test.go
+++ b/ably/realtime_conn_test.go
@@ -37,7 +37,7 @@ func TestRealtimeConn_Connect(t *testing.T) {
 	app, client := ablytest.NewRealtime(ably.ClientOptions{}.Listener(rec.Channel()))
 	defer safeclose(t, ablytest.FullRealtimeCloser(client), app)
 
-	if err := await(client.Connection.State, ably.StateConnConnected); err != nil {
+	if err := await(client.Connection.StateEnum, ably.StateConnConnected); err != nil {
 		t.Fatal(err)
 	}
 	if serial := client.Connection.Serial(); serial != -1 {
@@ -89,13 +89,13 @@ func TestRealtimeConn_ConnectClose(t *testing.T) {
 	app, client := ablytest.NewRealtime(ably.ClientOptions{}.Listener(rec.Channel()))
 	defer safeclose(t, ablytest.FullRealtimeCloser(client), app)
 
-	if err := await(client.Connection.State, ably.StateConnConnected); err != nil {
+	if err := await(client.Connection.StateEnum, ably.StateConnConnected); err != nil {
 		t.Fatal(err)
 	}
 	if err := ablytest.FullRealtimeCloser(client).Close(); err != nil {
 		t.Fatalf("ablytest.FullRealtimeCloser(client).Close()=%v", err)
 	}
-	if err := await(client.Connection.State, ably.StateConnClosed); err != nil {
+	if err := await(client.Connection.StateEnum, ably.StateConnClosed); err != nil {
 		t.Fatal(err)
 	}
 	rec.Stop()
@@ -134,7 +134,7 @@ func TestRealtimeConn_AuthError(t *testing.T) {
 	if err = ablytest.ConnWaiter(client, client.Connect, ably.ConnectionEventConnected).Wait(); err == nil {
 		t.Fatal("Connect(): want err != nil")
 	}
-	if state := client.Connection.State(); state != ably.StateConnFailed {
+	if state := client.Connection.State(); state != ably.ConnectionStateFailed {
 		t.Fatalf("want state=%s; got %s", ably.StateConnFailed, state)
 	}
 }