package codec import ( "bytes" "context" "errors" "fmt" "net" "sync" "time" "github.com/oarkflow/mq/consts" ) // MockConn implements net.Conn for testing type MockConn struct { ReadBuffer *bytes.Buffer WriteBuffer *bytes.Buffer ReadDelay time.Duration WriteDelay time.Duration IsClosed bool ReadErr error WriteErr error mu sync.Mutex } // NewMockConn creates a new mock connection func NewMockConn() *MockConn { return &MockConn{ ReadBuffer: bytes.NewBuffer(nil), WriteBuffer: bytes.NewBuffer(nil), } } // Read implements the net.Conn Read method func (m *MockConn) Read(b []byte) (n int, err error) { m.mu.Lock() defer m.mu.Unlock() if m.IsClosed { return 0, errors.New("connection closed") } if m.ReadErr != nil { return 0, m.ReadErr } if m.ReadDelay > 0 { time.Sleep(m.ReadDelay) } return m.ReadBuffer.Read(b) } // Write implements the net.Conn Write method func (m *MockConn) Write(b []byte) (n int, err error) { m.mu.Lock() defer m.mu.Unlock() if m.IsClosed { return 0, errors.New("connection closed") } if m.WriteErr != nil { return 0, m.WriteErr } if m.WriteDelay > 0 { time.Sleep(m.WriteDelay) } return m.WriteBuffer.Write(b) } // Close implements the net.Conn Close method func (m *MockConn) Close() error { m.mu.Lock() defer m.mu.Unlock() m.IsClosed = true return nil } // LocalAddr implements the net.Conn LocalAddr method func (m *MockConn) LocalAddr() net.Addr { return &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0} } // RemoteAddr implements the net.Conn RemoteAddr method func (m *MockConn) RemoteAddr() net.Addr { return &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0} } // SetDeadline implements the net.Conn SetDeadline method func (m *MockConn) SetDeadline(t time.Time) error { return nil } // SetReadDeadline implements the net.Conn SetReadDeadline method func (m *MockConn) SetReadDeadline(t time.Time) error { return nil } // SetWriteDeadline implements the net.Conn SetWriteDeadline method func (m *MockConn) SetWriteDeadline(t time.Time) error { return nil } // CodecTestSuite provides utilities for testing the codec type CodecTestSuite struct { Codec *Codec Config *Config } // NewCodecTestSuite creates a new codec test suite func NewCodecTestSuite() *CodecTestSuite { config := DefaultConfig() // Set smaller timeouts for testing config.ReadTimeout = 500 * time.Millisecond config.WriteTimeout = 500 * time.Millisecond return &CodecTestSuite{ Codec: NewCodec(config), Config: config, } } // SendReceiveTest tests sending and receiving a message func (ts *CodecTestSuite) SendReceiveTest(msg *Message) error { conn := NewMockConn() // Send the message ctx := context.Background() if err := ts.Codec.SendMessage(ctx, conn, msg); err != nil { return fmt.Errorf("failed to send message: %w", err) } // Move written data to read buffer to simulate network transport conn.ReadBuffer.Write(conn.WriteBuffer.Bytes()) conn.WriteBuffer.Reset() // Receive the message received, err := ts.Codec.ReadMessage(ctx, conn) if err != nil { return fmt.Errorf("failed to receive message: %w", err) } // Validate the message if received.Command != msg.Command { return fmt.Errorf("command mismatch: got %v, want %v", received.Command, msg.Command) } if received.Queue != msg.Queue { return fmt.Errorf("queue mismatch: got %v, want %v", received.Queue, msg.Queue) } if !bytes.Equal(received.Payload, msg.Payload) { return fmt.Errorf("payload mismatch: got %d bytes, want %d bytes", len(received.Payload), len(msg.Payload)) } return nil } // FragmentationTest tests the fragmentation and reassembly of large messages func (ts *CodecTestSuite) FragmentationTest(payload []byte) error { msg := &Message{ Command: consts.CMD(1), // Use appropriate command from your consts Queue: "test_queue", Headers: map[string]string{"test": "header"}, Payload: payload, Version: ProtocolVersion, Timestamp: time.Now().Unix(), ID: "test-message-id", } conn := NewMockConn() // Configure fragmentation fm := NewFragmentManager(ts.Codec, ts.Config) // Send the fragmented message ctx := context.Background() if err := fm.sendFragmentedMessage(ctx, conn, msg); err != nil { return fmt.Errorf("failed to send fragmented message: %w", err) } // Move written data to read buffer to simulate network transport conn.ReadBuffer.Write(conn.WriteBuffer.Bytes()) conn.WriteBuffer.Reset() // Receive and reassemble the message received, err := ts.Codec.ReadMessage(ctx, conn) if err != nil { return fmt.Errorf("failed to receive message: %w", err) } // Validate the reassembled message if received.Command != msg.Command { return fmt.Errorf("command mismatch: got %v, want %v", received.Command, msg.Command) } if received.Queue != msg.Queue { return fmt.Errorf("queue mismatch: got %v, want %v", received.Queue, msg.Queue) } if !bytes.Equal(received.Payload, msg.Payload) { return fmt.Errorf("payload mismatch: got %d bytes, want %d bytes", len(received.Payload), len(msg.Payload)) } return nil }