utils/network/udp_test.go

package network

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

const (
	testUDPServerAddress           = "127.0.0.1:10060"
	testUDPServerTimeout           = time.Second
	testUDPServerReceiveBufferSize = 1024

	testUDPClientTimeout           = time.Second
	testUDPClientReceiveBufferSize = 1024
)

func TestUDP(t *testing.T) {
	server := &UDPServer{}

	err := server.Connect(testUDPServerAddress, NewUDPServerOptions(
		WithUDPServerReadTimeout(testUDPServerTimeout),
		WithUDPServerWriteTimeout(testUDPServerTimeout),
		WithUDPServerReceiveBufferSize(testUDPServerReceiveBufferSize),
		WithUDPServerRequestCallback(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 server.Disconnect()

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

	client := &UDPClient{}
	err = client.Connect(testUDPServerAddress, NewUDPClientOptions(
		WithUDPClientReadTimeout(testUDPClientTimeout),
		WithUDPClientWriteTimeout(testUDPClientTimeout),
		WithUDPClientReceiveBufferSize(testUDPClientReceiveBufferSize),
		WithUDPClientRequestNonBlockCount(2),
		WithUDPClientResponseCallback(func(dataReader *DataReader) {
			requestBytes, err := dataReader.Bytes(dataReader.Len())
			if err != nil {
				t.Fatal(err)
				return
			}

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

	defer client.Disconnect()

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

	wg.Wait()
}