test: fix integration_tests for trafficlogger

core/internal/integration_tests/mocks/mock_TrafficLogger.go

@@ -2,7 +2,12 @@
 
 package mocks
 
-import mock "github.com/stretchr/testify/mock"
+import (
+	quic "github.com/apernet/quic-go"
+	mock "github.com/stretchr/testify/mock"
+
+	server "github.com/apernet/hysteria/core/v2/server"
+)
 
 // MockTrafficLogger is an autogenerated mock type for the TrafficLogger type
 type MockTrafficLogger struct {
@@ -99,6 +104,73 @@ func (_c *MockTrafficLogger_LogTraffic_Call) RunAndReturn(run func(string, uint6
 	return _c
 }
 
+// TraceStream provides a mock function with given fields: stream, stats
+func (_m *MockTrafficLogger) TraceStream(stream quic.Stream, stats *server.StreamStats) {
+	_m.Called(stream, stats)
+}
+
+// MockTrafficLogger_TraceStream_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'TraceStream'
+type MockTrafficLogger_TraceStream_Call struct {
+	*mock.Call
+}
+
+// TraceStream is a helper method to define mock.On call
+//   - stream quic.Stream
+//   - stats *server.StreamStats
+func (_e *MockTrafficLogger_Expecter) TraceStream(stream interface{}, stats interface{}) *MockTrafficLogger_TraceStream_Call {
+	return &MockTrafficLogger_TraceStream_Call{Call: _e.mock.On("TraceStream", stream, stats)}
+}
+
+func (_c *MockTrafficLogger_TraceStream_Call) Run(run func(stream quic.Stream, stats *server.StreamStats)) *MockTrafficLogger_TraceStream_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(quic.Stream), args[1].(*server.StreamStats))
+	})
+	return _c
+}
+
+func (_c *MockTrafficLogger_TraceStream_Call) Return() *MockTrafficLogger_TraceStream_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockTrafficLogger_TraceStream_Call) RunAndReturn(run func(quic.Stream, *server.StreamStats)) *MockTrafficLogger_TraceStream_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// UntraceStream provides a mock function with given fields: stream
+func (_m *MockTrafficLogger) UntraceStream(stream quic.Stream) {
+	_m.Called(stream)
+}
+
+// MockTrafficLogger_UntraceStream_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UntraceStream'
+type MockTrafficLogger_UntraceStream_Call struct {
+	*mock.Call
+}
+
+// UntraceStream is a helper method to define mock.On call
+//   - stream quic.Stream
+func (_e *MockTrafficLogger_Expecter) UntraceStream(stream interface{}) *MockTrafficLogger_UntraceStream_Call {
+	return &MockTrafficLogger_UntraceStream_Call{Call: _e.mock.On("UntraceStream", stream)}
+}
+
+func (_c *MockTrafficLogger_UntraceStream_Call) Run(run func(stream quic.Stream)) *MockTrafficLogger_UntraceStream_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(quic.Stream))
+	})
+	return _c
+}
+
+func (_c *MockTrafficLogger_UntraceStream_Call) Return() *MockTrafficLogger_UntraceStream_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockTrafficLogger_UntraceStream_Call) RunAndReturn(run func(quic.Stream)) *MockTrafficLogger_UntraceStream_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 {

core/internal/integration_tests/trafficlogger_test.go

@@ -62,6 +62,7 @@ func TestClientServerTrafficLoggerTCP(t *testing.T) {
 		return nil
 	})
 	serverOb.EXPECT().TCP(addr).Return(sobConn, nil).Once()
+	trafficLogger.EXPECT().TraceStream(mock.Anything, mock.Anything).Return().Once()
 
 	conn, err := c.TCP(addr)
 	assert.NoError(t, err)
@@ -84,6 +85,7 @@ func TestClientServerTrafficLoggerTCP(t *testing.T) {
 	time.Sleep(1 * time.Second) // Need some time for the server to receive the data
 
 	// Client reads from server again but blocked
+	trafficLogger.EXPECT().UntraceStream(mock.Anything).Return().Once()
 	trafficLogger.EXPECT().LogTraffic("nobody", uint64(0), uint64(4)).Return(false).Once()
 	trafficLogger.EXPECT().LogOnlineState("nobody", false).Return().Once()
 	sobConnCh <- []byte("nope")