github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/mux/mux_test.go (about)

     1  package mux
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  	"sync"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    12  	"github.com/Asutorufa/yuhaiin/pkg/net/proxy/simple"
    13  	"github.com/Asutorufa/yuhaiin/pkg/protos/node/protocol"
    14  	"github.com/Asutorufa/yuhaiin/pkg/utils/assert"
    15  )
    16  
    17  func TestMux(t *testing.T) {
    18  	lis, err := net.Listen("tcp", "127.0.0.1:4431")
    19  	assert.NoError(t, err)
    20  	defer lis.Close()
    21  
    22  	ms := newServer(lis)
    23  	defer ms.Close()
    24  
    25  	wg := sync.WaitGroup{}
    26  	wg.Add(1)
    27  	go func() {
    28  		defer wg.Done()
    29  
    30  		conn, err := ms.Accept()
    31  		assert.NoError(t, err)
    32  
    33  		data, err := io.ReadAll(conn)
    34  		assert.NoError(t, err)
    35  
    36  		t.Log(string(data))
    37  	}()
    38  
    39  	p, err := simple.NewClient(&protocol.Protocol_Simple{
    40  		Simple: &protocol.Simple{
    41  			Host: "127.0.0.1",
    42  			Port: 4431,
    43  		},
    44  	})(nil)
    45  	assert.NoError(t, err)
    46  
    47  	p, err = NewClient(&protocol.Protocol_Mux{
    48  		Mux: &protocol.Mux{
    49  			Concurrency: 1,
    50  		},
    51  	})(p)
    52  	assert.NoError(t, err)
    53  
    54  	conn, err := p.Conn(context.TODO(), netapi.EmptyAddr)
    55  	assert.NoError(t, err)
    56  	defer conn.Close()
    57  
    58  	_, err = conn.Write([]byte(time.Now().String()))
    59  	assert.NoError(t, err)
    60  
    61  	conn.Close()
    62  
    63  	wg.Wait()
    64  }