From cfce4b255e946bc67ac020717c6fbeff25da89ed Mon Sep 17 00:00:00 2001
From: sacOO7 <sachinshinde7676@gmail.com>
Date: Mon, 2 Sep 2024 13:18:10 +0530
Subject: [PATCH] Added integration test for server initiated auth, refactored
 related test name

---
 ably/ably_test.go                           |  3 ++
 ably/proto_protocol_message.go              |  3 ++
 ably/realtime_conn_spec_integration_test.go | 53 ++++++++++++++++++++-
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/ably/ably_test.go b/ably/ably_test.go
index 89fd0627..9b5cf2f2 100644
--- a/ably/ably_test.go
+++ b/ably/ably_test.go
@@ -272,6 +272,9 @@ func (rec *MessageRecorder) CheckIfReceived(action ably.ProtoAction, times int)
 				}
 			}
 		}
+		if times == 0 && times == counter {
+			return true
+		}
 		return false
 	}
 }
diff --git a/ably/proto_protocol_message.go b/ably/proto_protocol_message.go
index 4b6542c2..2fc460b6 100644
--- a/ably/proto_protocol_message.go
+++ b/ably/proto_protocol_message.go
@@ -183,6 +183,9 @@ func (msg *protocolMessage) String() string {
 	case actionMessage:
 		return fmt.Sprintf("(action=%q, id=%q, messages=%v)", msg.Action,
 			msg.ConnectionID, msg.Messages)
+	case actionAuth:
+		return fmt.Sprintf("(action=%q, id=%q, auth=%v)", msg.Action,
+			msg.ConnectionID, msg.Auth)
 	default:
 		return fmt.Sprintf("%#v", msg)
 	}
diff --git a/ably/realtime_conn_spec_integration_test.go b/ably/realtime_conn_spec_integration_test.go
index c78846bc..f5e8604a 100644
--- a/ably/realtime_conn_spec_integration_test.go
+++ b/ably/realtime_conn_spec_integration_test.go
@@ -1783,7 +1783,7 @@ func TestRealtimeConn_RTN15h3_Success(t *testing.T) {
 	ablytest.Instantly.NoRecv(t, nil, stateChanges, t.Fatalf)
 }
 
-func TestRealtimeConn_RTN15h_Integration_ClientInitiatedAuth(t *testing.T) {
+func TestRealtimeConn_RTN22a_RTN15h2_Integration_ServerInitiatedAuth(t *testing.T) {
 	t.Parallel()
 	app, restClient := ablytest.NewREST()
 	defer safeclose(t, app)
@@ -1833,6 +1833,57 @@ func TestRealtimeConn_RTN15h_Integration_ClientInitiatedAuth(t *testing.T) {
 	assert.ElementsMatch(t, authCallbackTokens, tokens)
 }
 
+func TestRealtimeConn_RTN22_RTC8_Integration_ServerInitiatedAuth(t *testing.T) {
+	app, restClient := ablytest.NewREST()
+	defer safeclose(t, app)
+
+	recorder := NewMessageRecorder()
+	authCallbackTokens := []string{}
+
+	// Server sends AUTH message 30 seconds before token expiry.
+	// So sending client token with expiry of 33 seconds, server will send AUTH msg after 3 seconds.
+	authCallback := func(ctx context.Context, tp ably.TokenParams) (ably.Tokener, error) {
+		tokenExpiry := 33000
+		token, err := restClient.Auth.RequestToken(context.Background(), &ably.TokenParams{TTL: int64(tokenExpiry)})
+		authCallbackTokens = append(authCallbackTokens, token.Token)
+		return token, err
+	}
+
+	realtime, err := ably.NewRealtime(
+		ably.WithAutoConnect(false),
+		ably.WithDial(recorder.Dial),
+		ably.WithUseBinaryProtocol(false),
+		ably.WithEnvironment(ablytest.Environment),
+		ably.WithAuthCallback(authCallback))
+
+	assert.NoError(t, err)
+	defer realtime.Close()
+
+	err = ablytest.Wait(ablytest.ConnWaiter(realtime, realtime.Connect, ably.ConnectionEventConnected), nil)
+	assert.NoError(t, err)
+
+	for i := 0; i < 3; i++ {
+		// auth msg sent by ably server every 3 seconds, so connection is updated by client
+		err = ablytest.Wait(ablytest.ConnWaiter(realtime, nil, ably.ConnectionEventUpdate), nil)
+		assert.NoError(t, err)
+		assert.Equal(t, ably.ConnectionStateConnected, realtime.Connection.State())
+		assert.True(t, ablytest.Instantly.IsTrue(recorder.CheckIfReceived(ably.ActionAuth, i+1)))
+	}
+	assert.True(t, ablytest.Instantly.IsTrue(recorder.CheckIfReceived(ably.ActionDisconnected, 0)))
+
+	// Only one dial attempt
+	tokens := []string{}
+	assert.Len(t, recorder.URLs(), 1)
+	for _, url := range recorder.URLs() {
+		tokens = append(tokens, url.Query().Get("access_token"))
+	}
+	assert.Len(t, tokens, 1)
+
+	assert.Len(t, authCallbackTokens, 4)
+	assert.Equal(t, tokens[0], authCallbackTokens[0])
+	assertUnique(t, authCallbackTokens)
+}
+
 func TestRealtimeConn_RTN15i_OnErrorWhenConnected(t *testing.T) {
 
 	in := make(chan *ably.ProtocolMessage, 1)