feat(wip): test reworks

core/client/udp_test.go

@@ -6,7 +6,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/magiconair/properties/assert"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"go.uber.org/goleak"
 
@@ -28,9 +28,9 @@ func TestUDPSessionManager(t *testing.T) {
 
 	// Test UDP session IO
 	udpConn1, err := sm.NewUDP()
-	assert.Equal(t, err, nil)
+	assert.NoError(t, err)
 	udpConn2, err := sm.NewUDP()
-	assert.Equal(t, err, nil)
+	assert.NoError(t, err)
 
 	msg1 := &protocol.UDPMessage{
 		SessionID: 1,
@@ -42,7 +42,7 @@ func TestUDPSessionManager(t *testing.T) {
 	}
 	io.EXPECT().SendMessage(mock.Anything, msg1).Return(nil).Once()
 	err = udpConn1.Send(msg1.Data, msg1.Addr)
-	assert.Equal(t, err, nil)
+	assert.NoError(t, err)
 
 	msg2 := &protocol.UDPMessage{
 		SessionID: 2,
@@ -54,7 +54,7 @@ func TestUDPSessionManager(t *testing.T) {
 	}
 	io.EXPECT().SendMessage(mock.Anything, msg2).Return(nil).Once()
 	err = udpConn2.Send(msg2.Data, msg2.Addr)
-	assert.Equal(t, err, nil)
+	assert.NoError(t, err)
 
 	respMsg1 := &protocol.UDPMessage{
 		SessionID: 1,
@@ -66,7 +66,7 @@ func TestUDPSessionManager(t *testing.T) {
 	}
 	receiveCh <- respMsg1
 	data, addr, err := udpConn1.Receive()
-	assert.Equal(t, err, nil)
+	assert.NoError(t, err)
 	assert.Equal(t, data, respMsg1.Data)
 	assert.Equal(t, addr, respMsg1.Addr)
 
@@ -80,7 +80,7 @@ func TestUDPSessionManager(t *testing.T) {
 	}
 	receiveCh <- respMsg2
 	data, addr, err = udpConn2.Receive()
-	assert.Equal(t, err, nil)
+	assert.NoError(t, err)
 	assert.Equal(t, data, respMsg2.Data)
 	assert.Equal(t, addr, respMsg2.Addr)
 
@@ -101,7 +101,7 @@ func TestUDPSessionManager(t *testing.T) {
 		_, _, err := udpConn1.Receive()
 		errChan <- err
 	}()
-	assert.Equal(t, udpConn1.Close(), nil)
+	assert.NoError(t, udpConn1.Close())
 	assert.Equal(t, <-errChan, io2.EOF)
 
 	// Test close IO unblocks Receive() and blocks new UDP creation
@@ -117,6 +117,6 @@ func TestUDPSessionManager(t *testing.T) {
 
 	// Leak checks
 	time.Sleep(1 * time.Second)
-	assert.Equal(t, sm.Count(), 0, "session count should be 0")
+	assert.Zero(t, sm.Count(), "session count should be 0")
 	goleak.VerifyNone(t)
 }

core/internal/integration_tests/.mockery.yaml

@@ -0,0 +1,15 @@
+with-expecter: true
+dir: mocks
+outpkg: mocks
+packages:
+  github.com/apernet/hysteria/core/server:
+    interfaces:
+      Authenticator:
+        config:
+          mockname: MockAuthenticator
+      EventLogger:
+        config:
+          mockname: MockEventLogger
+      TrafficLogger:
+        config:
+          mockname: MockTrafficLogger

core/internal/integration_tests/close_test.go

@@ -1,13 +1,16 @@
 package integration_tests
 
 import (
-	"bytes"
 	"crypto/rand"
 	"io"
 	"net"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
 	"github.com/apernet/hysteria/core/client"
+	"github.com/apernet/hysteria/core/internal/integration_tests/mocks"
 	"github.com/apernet/hysteria/core/server"
 )
 
@@ -17,41 +20,31 @@ func TestClientServerTCPClose(t *testing.T) {
 	// Create server
 	udpAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 14514}
 	udpConn, err := net.ListenUDP("udp", udpAddr)
-	if err != nil {
-		t.Fatal("error creating server:", err)
-	}
+	assert.NoError(t, err)
+	auth := mocks.NewMockAuthenticator(t)
+	auth.EXPECT().Authenticate(mock.Anything, mock.Anything, mock.Anything).Return(true, "nobody")
 	s, err := server.NewServer(&server.Config{
-		TLSConfig: serverTLSConfig(),
-		Conn:      udpConn,
-		Authenticator: &pwAuthenticator{
-			Password: "password",
-			ID:       "nobody",
-		},
+		TLSConfig:     serverTLSConfig(),
+		Conn:          udpConn,
+		Authenticator: auth,
 	})
-	if err != nil {
-		t.Fatal("error creating server:", err)
-	}
+	assert.NoError(t, err)
 	defer s.Close()
 	go s.Serve()
 
 	// Create client
 	c, err := client.NewClient(&client.Config{
 		ServerAddr: udpAddr,
-		Auth:       "password",
 		TLSConfig:  client.TLSConfig{InsecureSkipVerify: true},
 	})
-	if err != nil {
-		t.Fatal("error creating client:", err)
-	}
+	assert.NoError(t, err)
 	defer c.Close()
 
 	t.Run("Close local", func(t *testing.T) {
 		// TCP sink server
 		sinkAddr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 33344}
 		sinkListener, err := net.ListenTCP("tcp", sinkAddr)
-		if err != nil {
-			t.Fatal("error creating sink server:", err)
-		}
+		assert.NoError(t, err)
 		sinkCh := make(chan sinkEvent, 1)
 		sinkServer := &tcpSinkServer{
 			Listener: sinkListener,
@@ -63,47 +56,33 @@ func TestClientServerTCPClose(t *testing.T) {
 		// Generate some random data
 		sData := make([]byte, 1024000)
 		_, err = rand.Read(sData)
-		if err != nil {
-			t.Fatal("error generating random data:", err)
-		}
+		assert.NoError(t, err)
 
 		// Dial and send data to TCP sink server
 		conn, err := c.DialTCP(sinkAddr.String())
-		if err != nil {
-			t.Fatal("error dialing TCP:", err)
-		}
-		defer conn.Close()
+		assert.NoError(t, err)
 		_, err = conn.Write(sData)
-		if err != nil {
-			t.Fatal("error writing to TCP:", err)
-		}
+		assert.NoError(t, err)
 
 		// Close the connection
 		// This should cause the sink server to send an event to the channel
 		_ = conn.Close()
 		event := <-sinkCh
-		if event.Err != nil {
-			t.Fatal("non-nil error received from sink server:", event.Err)
-		}
-		if !bytes.Equal(event.Data, sData) {
-			t.Fatal("data mismatch")
-		}
+		assert.NoError(t, event.Err)
+		assert.Equal(t, sData, event.Data)
 	})
 
 	t.Run("Close remote", func(t *testing.T) {
 		// Generate some random data
 		sData := make([]byte, 1024000)
 		_, err = rand.Read(sData)
-		if err != nil {
-			t.Fatal("error generating random data:", err)
-		}
+		assert.NoError(t, err)
 
 		// TCP sender server
 		senderAddr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 33345}
 		senderListener, err := net.ListenTCP("tcp", senderAddr)
-		if err != nil {
-			t.Fatal("error creating sender server:", err)
-		}
+		assert.NoError(t, err)
+
 		senderServer := &tcpSenderServer{
 			Listener: senderListener,
 			Data:     sData,
@@ -113,17 +92,11 @@ func TestClientServerTCPClose(t *testing.T) {
 
 		// Dial and read data from TCP sender server
 		conn, err := c.DialTCP(senderAddr.String())
-		if err != nil {
-			t.Fatal("error dialing TCP:", err)
-		}
+		assert.NoError(t, err)
 		defer conn.Close()
 		rData, err := io.ReadAll(conn)
-		if err != nil {
-			t.Fatal("error reading from TCP:", err)
-		}
-		if !bytes.Equal(rData, sData) {
-			t.Fatal("data mismatch")
-		}
+		assert.NoError(t, err)
+		assert.Equal(t, sData, rData)
 	})
 }
 

core/internal/integration_tests/mocks/mock_Authenticator.go

@@ -0,0 +1,90 @@
+// Code generated by mockery v2.32.0. DO NOT EDIT.
+
+package mocks
+
+import (
+	net "net"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// MockAuthenticator is an autogenerated mock type for the Authenticator type
+type MockAuthenticator struct {
+	mock.Mock
+}
+
+type MockAuthenticator_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockAuthenticator) EXPECT() *MockAuthenticator_Expecter {
+	return &MockAuthenticator_Expecter{mock: &_m.Mock}
+}
+
+// Authenticate provides a mock function with given fields: addr, auth, tx
+func (_m *MockAuthenticator) Authenticate(addr net.Addr, auth string, tx uint64) (bool, string) {
+	ret := _m.Called(addr, auth, tx)
+
+	var r0 bool
+	var r1 string
+	if rf, ok := ret.Get(0).(func(net.Addr, string, uint64) (bool, string)); ok {
+		return rf(addr, auth, tx)
+	}
+	if rf, ok := ret.Get(0).(func(net.Addr, string, uint64) bool); ok {
+		r0 = rf(addr, auth, tx)
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	if rf, ok := ret.Get(1).(func(net.Addr, string, uint64) string); ok {
+		r1 = rf(addr, auth, tx)
+	} else {
+		r1 = ret.Get(1).(string)
+	}
+
+	return r0, r1
+}
+
+// MockAuthenticator_Authenticate_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Authenticate'
+type MockAuthenticator_Authenticate_Call struct {
+	*mock.Call
+}
+
+// Authenticate is a helper method to define mock.On call
+//   - addr net.Addr
+//   - auth string
+//   - tx uint64
+func (_e *MockAuthenticator_Expecter) Authenticate(addr interface{}, auth interface{}, tx interface{}) *MockAuthenticator_Authenticate_Call {
+	return &MockAuthenticator_Authenticate_Call{Call: _e.mock.On("Authenticate", addr, auth, tx)}
+}
+
+func (_c *MockAuthenticator_Authenticate_Call) Run(run func(addr net.Addr, auth string, tx uint64)) *MockAuthenticator_Authenticate_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(net.Addr), args[1].(string), args[2].(uint64))
+	})
+	return _c
+}
+
+func (_c *MockAuthenticator_Authenticate_Call) Return(ok bool, id string) *MockAuthenticator_Authenticate_Call {
+	_c.Call.Return(ok, id)
+	return _c
+}
+
+func (_c *MockAuthenticator_Authenticate_Call) RunAndReturn(run func(net.Addr, string, uint64) (bool, string)) *MockAuthenticator_Authenticate_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockAuthenticator creates a new instance of MockAuthenticator. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockAuthenticator(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockAuthenticator {
+	mock := &MockAuthenticator{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

core/internal/integration_tests/mocks/mock_EventLogger.go

@@ -0,0 +1,249 @@
+// Code generated by mockery v2.32.0. DO NOT EDIT.
+
+package mocks
+
+import (
+	net "net"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// MockEventLogger is an autogenerated mock type for the EventLogger type
+type MockEventLogger struct {
+	mock.Mock
+}
+
+type MockEventLogger_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockEventLogger) EXPECT() *MockEventLogger_Expecter {
+	return &MockEventLogger_Expecter{mock: &_m.Mock}
+}
+
+// Connect provides a mock function with given fields: addr, id, tx
+func (_m *MockEventLogger) Connect(addr net.Addr, id string, tx uint64) {
+	_m.Called(addr, id, tx)
+}
+
+// MockEventLogger_Connect_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Connect'
+type MockEventLogger_Connect_Call struct {
+	*mock.Call
+}
+
+// Connect is a helper method to define mock.On call
+//   - addr net.Addr
+//   - id string
+//   - tx uint64
+func (_e *MockEventLogger_Expecter) Connect(addr interface{}, id interface{}, tx interface{}) *MockEventLogger_Connect_Call {
+	return &MockEventLogger_Connect_Call{Call: _e.mock.On("Connect", addr, id, tx)}
+}
+
+func (_c *MockEventLogger_Connect_Call) Run(run func(addr net.Addr, id string, tx uint64)) *MockEventLogger_Connect_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(net.Addr), args[1].(string), args[2].(uint64))
+	})
+	return _c
+}
+
+func (_c *MockEventLogger_Connect_Call) Return() *MockEventLogger_Connect_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockEventLogger_Connect_Call) RunAndReturn(run func(net.Addr, string, uint64)) *MockEventLogger_Connect_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Disconnect provides a mock function with given fields: addr, id, err
+func (_m *MockEventLogger) Disconnect(addr net.Addr, id string, err error) {
+	_m.Called(addr, id, err)
+}
+
+// MockEventLogger_Disconnect_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Disconnect'
+type MockEventLogger_Disconnect_Call struct {
+	*mock.Call
+}
+
+// Disconnect is a helper method to define mock.On call
+//   - addr net.Addr
+//   - id string
+//   - err error
+func (_e *MockEventLogger_Expecter) Disconnect(addr interface{}, id interface{}, err interface{}) *MockEventLogger_Disconnect_Call {
+	return &MockEventLogger_Disconnect_Call{Call: _e.mock.On("Disconnect", addr, id, err)}
+}
+
+func (_c *MockEventLogger_Disconnect_Call) Run(run func(addr net.Addr, id string, err error)) *MockEventLogger_Disconnect_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(net.Addr), args[1].(string), args[2].(error))
+	})
+	return _c
+}
+
+func (_c *MockEventLogger_Disconnect_Call) Return() *MockEventLogger_Disconnect_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockEventLogger_Disconnect_Call) RunAndReturn(run func(net.Addr, string, error)) *MockEventLogger_Disconnect_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// TCPError provides a mock function with given fields: addr, id, reqAddr, err
+func (_m *MockEventLogger) TCPError(addr net.Addr, id string, reqAddr string, err error) {
+	_m.Called(addr, id, reqAddr, err)
+}
+
+// MockEventLogger_TCPError_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'TCPError'
+type MockEventLogger_TCPError_Call struct {
+	*mock.Call
+}
+
+// TCPError is a helper method to define mock.On call
+//   - addr net.Addr
+//   - id string
+//   - reqAddr string
+//   - err error
+func (_e *MockEventLogger_Expecter) TCPError(addr interface{}, id interface{}, reqAddr interface{}, err interface{}) *MockEventLogger_TCPError_Call {
+	return &MockEventLogger_TCPError_Call{Call: _e.mock.On("TCPError", addr, id, reqAddr, err)}
+}
+
+func (_c *MockEventLogger_TCPError_Call) Run(run func(addr net.Addr, id string, reqAddr string, err error)) *MockEventLogger_TCPError_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(net.Addr), args[1].(string), args[2].(string), args[3].(error))
+	})
+	return _c
+}
+
+func (_c *MockEventLogger_TCPError_Call) Return() *MockEventLogger_TCPError_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockEventLogger_TCPError_Call) RunAndReturn(run func(net.Addr, string, string, error)) *MockEventLogger_TCPError_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// TCPRequest provides a mock function with given fields: addr, id, reqAddr
+func (_m *MockEventLogger) TCPRequest(addr net.Addr, id string, reqAddr string) {
+	_m.Called(addr, id, reqAddr)
+}
+
+// MockEventLogger_TCPRequest_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'TCPRequest'
+type MockEventLogger_TCPRequest_Call struct {
+	*mock.Call
+}
+
+// TCPRequest is a helper method to define mock.On call
+//   - addr net.Addr
+//   - id string
+//   - reqAddr string
+func (_e *MockEventLogger_Expecter) TCPRequest(addr interface{}, id interface{}, reqAddr interface{}) *MockEventLogger_TCPRequest_Call {
+	return &MockEventLogger_TCPRequest_Call{Call: _e.mock.On("TCPRequest", addr, id, reqAddr)}
+}
+
+func (_c *MockEventLogger_TCPRequest_Call) Run(run func(addr net.Addr, id string, reqAddr string)) *MockEventLogger_TCPRequest_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(net.Addr), args[1].(string), args[2].(string))
+	})
+	return _c
+}
+
+func (_c *MockEventLogger_TCPRequest_Call) Return() *MockEventLogger_TCPRequest_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockEventLogger_TCPRequest_Call) RunAndReturn(run func(net.Addr, string, string)) *MockEventLogger_TCPRequest_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// UDPError provides a mock function with given fields: addr, id, sessionID, err
+func (_m *MockEventLogger) UDPError(addr net.Addr, id string, sessionID uint32, err error) {
+	_m.Called(addr, id, sessionID, err)
+}
+
+// MockEventLogger_UDPError_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UDPError'
+type MockEventLogger_UDPError_Call struct {
+	*mock.Call
+}
+
+// UDPError is a helper method to define mock.On call
+//   - addr net.Addr
+//   - id string
+//   - sessionID uint32
+//   - err error
+func (_e *MockEventLogger_Expecter) UDPError(addr interface{}, id interface{}, sessionID interface{}, err interface{}) *MockEventLogger_UDPError_Call {
+	return &MockEventLogger_UDPError_Call{Call: _e.mock.On("UDPError", addr, id, sessionID, err)}
+}
+
+func (_c *MockEventLogger_UDPError_Call) Run(run func(addr net.Addr, id string, sessionID uint32, err error)) *MockEventLogger_UDPError_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(net.Addr), args[1].(string), args[2].(uint32), args[3].(error))
+	})
+	return _c
+}
+
+func (_c *MockEventLogger_UDPError_Call) Return() *MockEventLogger_UDPError_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockEventLogger_UDPError_Call) RunAndReturn(run func(net.Addr, string, uint32, error)) *MockEventLogger_UDPError_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// UDPRequest provides a mock function with given fields: addr, id, sessionID, reqAddr
+func (_m *MockEventLogger) UDPRequest(addr net.Addr, id string, sessionID uint32, reqAddr string) {
+	_m.Called(addr, id, sessionID, reqAddr)
+}
+
+// MockEventLogger_UDPRequest_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UDPRequest'
+type MockEventLogger_UDPRequest_Call struct {
+	*mock.Call
+}
+
+// UDPRequest is a helper method to define mock.On call
+//   - addr net.Addr
+//   - id string
+//   - sessionID uint32
+//   - reqAddr string
+func (_e *MockEventLogger_Expecter) UDPRequest(addr interface{}, id interface{}, sessionID interface{}, reqAddr interface{}) *MockEventLogger_UDPRequest_Call {
+	return &MockEventLogger_UDPRequest_Call{Call: _e.mock.On("UDPRequest", addr, id, sessionID, reqAddr)}
+}
+
+func (_c *MockEventLogger_UDPRequest_Call) Run(run func(addr net.Addr, id string, sessionID uint32, reqAddr string)) *MockEventLogger_UDPRequest_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(net.Addr), args[1].(string), args[2].(uint32), args[3].(string))
+	})
+	return _c
+}
+
+func (_c *MockEventLogger_UDPRequest_Call) Return() *MockEventLogger_UDPRequest_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockEventLogger_UDPRequest_Call) RunAndReturn(run func(net.Addr, string, uint32, string)) *MockEventLogger_UDPRequest_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockEventLogger creates a new instance of MockEventLogger. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockEventLogger(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockEventLogger {
+	mock := &MockEventLogger{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

core/internal/integration_tests/mocks/mock_TrafficLogger.go

@@ -0,0 +1,76 @@
+// Code generated by mockery v2.32.0. DO NOT EDIT.
+
+package mocks
+
+import mock "github.com/stretchr/testify/mock"
+
+// MockTrafficLogger is an autogenerated mock type for the TrafficLogger type
+type MockTrafficLogger struct {
+	mock.Mock
+}
+
+type MockTrafficLogger_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockTrafficLogger) EXPECT() *MockTrafficLogger_Expecter {
+	return &MockTrafficLogger_Expecter{mock: &_m.Mock}
+}
+
+// Log provides a mock function with given fields: id, tx, rx
+func (_m *MockTrafficLogger) Log(id string, tx uint64, rx uint64) bool {
+	ret := _m.Called(id, tx, rx)
+
+	var r0 bool
+	if rf, ok := ret.Get(0).(func(string, uint64, uint64) bool); ok {
+		r0 = rf(id, tx, rx)
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	return r0
+}
+
+// MockTrafficLogger_Log_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Log'
+type MockTrafficLogger_Log_Call struct {
+	*mock.Call
+}
+
+// Log is a helper method to define mock.On call
+//   - id string
+//   - tx uint64
+//   - rx uint64
+func (_e *MockTrafficLogger_Expecter) Log(id interface{}, tx interface{}, rx interface{}) *MockTrafficLogger_Log_Call {
+	return &MockTrafficLogger_Log_Call{Call: _e.mock.On("Log", id, tx, rx)}
+}
+
+func (_c *MockTrafficLogger_Log_Call) Run(run func(id string, tx uint64, rx uint64)) *MockTrafficLogger_Log_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(string), args[1].(uint64), args[2].(uint64))
+	})
+	return _c
+}
+
+func (_c *MockTrafficLogger_Log_Call) Return(ok bool) *MockTrafficLogger_Log_Call {
+	_c.Call.Return(ok)
+	return _c
+}
+
+func (_c *MockTrafficLogger_Log_Call) RunAndReturn(run func(string, uint64, uint64) bool) *MockTrafficLogger_Log_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockTrafficLogger creates a new instance of MockTrafficLogger. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockTrafficLogger(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockTrafficLogger {
+	mock := &MockTrafficLogger{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

core/server/udp_test.go

@@ -5,7 +5,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/magiconair/properties/assert"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"go.uber.org/goleak"