From a77a6f85a61c9e0e19877b1cbdaebfb50e00b9ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=A3=8E=E6=89=87=E6=BB=91=E7=BF=94=E7=BF=BC?=
 <Fangliding.fshxy@outlook.com>
Date: Tue, 8 Jul 2025 22:02:53 +0800
Subject: [PATCH] Update client_test.go

---
 common/mux/client_test.go | 135 ++++++++++++++++++++++++++++----------
 1 file changed, 100 insertions(+), 35 deletions(-)

diff --git a/common/mux/client_test.go b/common/mux/client_test.go
index 9b3e8719..9626e2a2 100644
--- a/common/mux/client_test.go
+++ b/common/mux/client_test.go
@@ -1,51 +1,116 @@
 package mux_test
 
 import (
+	"context"
 	"testing"
+	"time"
 
-	. "github.com/xtls/xray-core/common/mux"
+	"github.com/golang/mock/gomock"
+	"github.com/xtls/xray-core/common"
+	"github.com/xtls/xray-core/common/errors"
+	"github.com/xtls/xray-core/common/mux"
+	"github.com/xtls/xray-core/common/net"
+	"github.com/xtls/xray-core/common/session"
+	"github.com/xtls/xray-core/testing/mocks"
+	"github.com/xtls/xray-core/transport"
+	"github.com/xtls/xray-core/transport/pipe"
 )
 
-func TestSessionManagerAdd(t *testing.T) {
-	m := NewSessionManager()
+func TestIncrementalPickerFailure(t *testing.T) {
+	mockCtl := gomock.NewController(t)
+	defer mockCtl.Finish()
 
-	s := m.Allocate(0)
-	if s.ID != 1 {
-		t.Error("id: ", s.ID)
-	}
-	if m.Size() != 1 {
-		t.Error("size: ", m.Size())
+	mockWorkerFactory := mocks.NewMuxClientWorkerFactory(mockCtl)
+	mockWorkerFactory.EXPECT().Create().Return(nil, errors.New("test"))
+
+	picker := mux.IncrementalWorkerPicker{
+		Factory: mockWorkerFactory,
 	}
 
-	s = m.Allocate(0)
-	if s.ID != 2 {
-		t.Error("id: ", s.ID)
-	}
-	if m.Size() != 2 {
-		t.Error("size: ", m.Size())
-	}
-
-	s = &Session{
-		ID: 4,
-	}
-	m.Add(s)
-	if s.ID != 4 {
-		t.Error("id: ", s.ID)
-	}
-	if m.Size() != 3 {
-		t.Error("size: ", m.Size())
+	_, err := picker.PickAvailable()
+	if err == nil {
+		t.Error("expected error, but nil")
 	}
 }
 
-func TestSessionManagerClose(t *testing.T) {
-	m := NewSessionManager()
-	s := m.Allocate(0)
+func TestClientWorkerEOF(t *testing.T) {
+	reader, writer := pipe.New(pipe.WithoutSizeLimit())
+	common.Must(writer.Close())
 
-	if m.CloseIfNoSession() {
-		t.Error("able to close")
-	}
-	m.Remove(false, s.ID)
-	if !m.CloseIfNoSession() {
-		t.Error("not able to close")
+	worker, err := mux.NewClientWorker(transport.Link{Reader: reader, Writer: writer}, mux.ClientStrategy{})
+	common.Must(err)
+
+	time.Sleep(time.Millisecond * 500)
+
+	f := worker.Dispatch(context.Background(), nil)
+	if f {
+		t.Error("expected failed dispatching, but actually not")
 	}
 }
+
+func TestClientWorkerClose(t *testing.T) {
+	mockCtl := gomock.NewController(t)
+	defer mockCtl.Finish()
+
+	r1, w1 := pipe.New(pipe.WithoutSizeLimit())
+	worker1, err := mux.NewClientWorker(transport.Link{
+		Reader: r1,
+		Writer: w1,
+	}, mux.ClientStrategy{
+		MaxConcurrency: 4,
+		MaxConnection:  4,
+	})
+	common.Must(err)
+
+	r2, w2 := pipe.New(pipe.WithoutSizeLimit())
+	worker2, err := mux.NewClientWorker(transport.Link{
+		Reader: r2,
+		Writer: w2,
+	}, mux.ClientStrategy{
+		MaxConcurrency: 4,
+		MaxConnection:  4,
+	})
+	common.Must(err)
+
+	factory := mocks.NewMuxClientWorkerFactory(mockCtl)
+	gomock.InOrder(
+		factory.EXPECT().Create().Return(worker1, nil),
+		factory.EXPECT().Create().Return(worker2, nil),
+	)
+
+	picker := &mux.IncrementalWorkerPicker{
+		Factory: factory,
+	}
+	manager := &mux.ClientManager{
+		Picker: picker,
+	}
+
+	tr1, tw1 := pipe.New(pipe.WithoutSizeLimit())
+	ctx1 := session.ContextWithOutbounds(context.Background(), []*session.Outbound{{
+		Target: net.TCPDestination(net.DomainAddress("www.example.com"), 80),
+	}})
+	common.Must(manager.Dispatch(ctx1, &transport.Link{
+		Reader: tr1,
+		Writer: tw1,
+	}))
+	defer tw1.Close()
+
+	common.Must(w1.Close())
+
+	time.Sleep(time.Millisecond * 500)
+	if !worker1.Closed() {
+		t.Error("worker1 is not finished")
+	}
+
+	tr2, tw2 := pipe.New(pipe.WithoutSizeLimit())
+	ctx2 := session.ContextWithOutbounds(context.Background(), []*session.Outbound{{
+		Target: net.TCPDestination(net.DomainAddress("www.example.com"), 80),
+	}})
+	common.Must(manager.Dispatch(ctx2, &transport.Link{
+		Reader: tr2,
+		Writer: tw2,
+	}))
+	defer tw2.Close()
+
+	common.Must(w2.Close())
+}