Shaun Campbell
865f7c26c9
All checks were successful
Build and Push Image / docker-build-and-push (push) Successful in 2m28s
77 lines
2.1 KiB
Go
77 lines
2.1 KiB
Go
package pgwire
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/binary"
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestStartupMessageReadsDatabase(t *testing.T) {
|
|
raw := startupBytes("analytics", "ruckstack")
|
|
message, err := ReadStartupMessage(bytes.NewReader(raw))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if message.Database != "analytics" {
|
|
t.Fatalf("Database = %q", message.Database)
|
|
}
|
|
}
|
|
|
|
func TestReadStartupMessageRejectsMissingDatabase(t *testing.T) {
|
|
_, err := ReadStartupMessage(bytes.NewReader(startupBytes("", "ruckstack")))
|
|
if !errors.Is(err, ErrMissingDatabase) {
|
|
t.Errorf("ReadStartupMessage() error = %v, want ErrMissingDatabase", err)
|
|
}
|
|
}
|
|
|
|
func TestReplayForwardsStartupAndUnderlyingStream(t *testing.T) {
|
|
reader, writer := io.Pipe()
|
|
defer reader.Close()
|
|
go func() {
|
|
_, _ = writer.Write([]byte("tail"))
|
|
_ = writer.Close()
|
|
}()
|
|
connection := &readOnlyConn{Reader: reader}
|
|
replayed := Replay(connection, []byte("startup"))
|
|
got, err := io.ReadAll(replayed)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if string(got) != "startuptail" {
|
|
t.Errorf("Replay() = %q", got)
|
|
}
|
|
}
|
|
|
|
func startupBytes(database, user string) []byte {
|
|
body := make([]byte, 4)
|
|
binary.BigEndian.PutUint32(body, startupProtocolVersion)
|
|
if database != "" {
|
|
body = append(body, "database"...)
|
|
body = append(body, 0)
|
|
body = append(body, database...)
|
|
body = append(body, 0)
|
|
}
|
|
body = append(body, "user"...)
|
|
body = append(body, 0)
|
|
body = append(body, user...)
|
|
body = append(body, 0, 0)
|
|
message := make([]byte, 4, 4+len(body))
|
|
message = append(message, body...)
|
|
binary.BigEndian.PutUint32(message, uint32(len(message)))
|
|
return message
|
|
}
|
|
|
|
type readOnlyConn struct{ io.Reader }
|
|
|
|
func (c *readOnlyConn) Write([]byte) (int, error) { return 0, io.ErrClosedPipe }
|
|
func (c *readOnlyConn) Close() error { return nil }
|
|
func (c *readOnlyConn) LocalAddr() net.Addr { return nil }
|
|
func (c *readOnlyConn) RemoteAddr() net.Addr { return nil }
|
|
func (c *readOnlyConn) SetDeadline(time.Time) error { return nil }
|
|
func (c *readOnlyConn) SetReadDeadline(time.Time) error { return nil }
|
|
func (c *readOnlyConn) SetWriteDeadline(time.Time) error { return nil }
|