utils/network/tcp_test.go

package network

import (
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"
)

const (
	testTCPServerAddress           = "127.0.0.1:10060"
	testTCPServerTimeout           = time.Second
	testTCPServerReceiveBufferSize = 1024

	testTCPClientTimeout           = time.Second
	testTCPClientReceiveBufferSize = 1024
)

func TestTCP(t *testing.T) {
	server := &TCPServer{}

	err := server.Connect(testTCPServerAddress, NewTCPServerOptions(
		WithTCPServerReadTimeout(testTCPServerTimeout),
		WithTCPServerWriteTimeout(testTCPServerTimeout),
		WithTCPServerReceiveBufferSize(testTCPServerReceiveBufferSize),
		WithTCPServerRequestCallback(func(dataReader *DataReader) *DataReader {
			requestBytes, err := dataReader.Bytes(dataReader.Len())
			if err != nil {
				t.Fatal(err)
				return nil
			}

			responseWriter := NewDataWriter()
			err = responseWriter.Bytes([]byte(strings.ToUpper(string(requestBytes))))
			if err != nil {
				t.Fatal(err)
				return nil
			}

			return responseWriter.ToReader()
		}),
	))
	if err != nil {
		t.Fatal(err)
	}

	defer func(server *TCPServer) {
		err := server.Disconnect()
		if err != nil {
			t.Fatal(err)
		}
	}(server)

	wg := &sync.WaitGroup{}
	wg.Add(2)

	client := &TCPClient{}
	err = client.Connect(testTCPServerAddress, NewTCPClientOptions(
		WithTCPClientReadTimeout(testTCPClientTimeout),
		WithTCPClientWriteTimeout(testTCPClientTimeout),
		WithTCPClientReceiveBufferSize(testTCPClientReceiveBufferSize),
		WithTCPClientRequestNonBlockCount(2),
		WithTCPClientResponseCallback(func(dataReader *DataReader) bool {
			defer wg.Done()

			requestBytes, err := dataReader.Bytes(dataReader.Len())
			if err != nil {
				t.Fatal(err)
				return true
			}

			fmt.Println(string(requestBytes))
			return true
		}),
	))
	if err != nil {
		t.Fatal(err)
	}

	defer client.Disconnect()

	client.Send([]byte("test1"))
	client.Send([]byte("test2"))

	wg.Wait()
}