fix fix test

topfreegames · Mar 4, 2024 · a54b1f4 · a54b1f4
1 parent d313780
commit a54b1f4
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/agent/agent.go b/agent/agent.go
@@ -388,7 +388,12 @@ func (a *agentImpl) Kick(ctx context.Context) error {
 		ctx:  ctx,
 		data: p,
 	}
-	a.chSend <- pWrite
+
+	select {
+	case a.chSend <- pWrite:
+	case <-a.chDie:
+	}
+
 	return nil
 }
 

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -130,21 +130,29 @@ func TestKick(t *testing.T) {
 	dieChan := make(chan bool)
 	hbTime := time.Second
 
+	expectedBytes := []byte("kick")
+
 	mockConn := mocks.NewMockPlayerConn(ctrl)
-	mockEncoder.EXPECT().Encode(gomock.Any(), gomock.Nil()).Do(
-		func(typ packet.Type, d []byte) {
-			assert.EqualValues(t, packet.Kick, typ)
-		})
-	mockConn.EXPECT().Write(gomock.Any()).Return(0, nil)
+	mockEncoder.EXPECT().Encode(packet.Type(packet.Kick), gomock.Nil()).Return(expectedBytes, nil)
 	messageEncoder := message.NewMessagesEncoder(false)
 
 	mockSerializer.EXPECT().GetName()
 
+	ctx := context.Background()
+
+	exceptedWrite := pendingWrite{
+		ctx:  ctx,
+		data: expectedBytes,
+	}
+
 	sessionPool := session.NewSessionPool()
-	ag := newAgent(mockConn, mockDecoder, mockEncoder, mockSerializer, hbTime, 10, dieChan, messageEncoder, nil, sessionPool)
-	c := context.Background()
-	err := ag.Kick(c)
+	ag := newAgent(mockConn, mockDecoder, mockEncoder, mockSerializer, hbTime, 10, dieChan, messageEncoder, nil, sessionPool).(*agentImpl)
+
+	err := ag.Kick(ctx)
 	assert.NoError(t, err)
+
+	recvData := helpers.ShouldEventuallyReceive(t, ag.chSend).(pendingWrite)
+	assert.Equal(t, exceptedWrite, recvData)
 }
 
 func TestAgentSend(t *testing.T) {