From cf225c83656f5bf74a0a4364f23ec1ed27b68298 Mon Sep 17 00:00:00 2001
From: Kris Wehner <kris@yelp.com>
Date: Tue, 3 Jan 2017 13:56:34 -0800
Subject: [PATCH] Add the ability to set all the fields in the constructor

---
 replication.go      | 28 +++++++++++++++++++++++-----
 replication_test.go | 27 ++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/replication.go b/replication.go
index fdbdce92..7d4c56e2 100644
--- a/replication.go
+++ b/replication.go
@@ -124,11 +124,29 @@ type StandbyStatus struct {
 
 // Create a standby status struct, which sets all the WAL positions
 // to the given wal position, and the client time to the current time.
-func NewStandbyStatus(walPosition uint64) (status *StandbyStatus) {
-	status = new(StandbyStatus)
-	status.WalFlushPosition = walPosition
-	status.WalApplyPosition = walPosition
-	status.WalWritePosition = walPosition
+// The wal positions are, in order:
+// WalFlushPosition
+// WalApplyPosition
+// WalWritePosition
+//
+// If only one position is provided, it will be used as the value for all 3
+// status fields. Note you must provide either 1 wal position, or all 3
+// in order to initialize the standby status.
+func NewStandbyStatus(walPositions ...uint64) (status *StandbyStatus, err error) {
+	if len(walPositions) == 1 {
+		status = new(StandbyStatus)
+		status.WalFlushPosition = walPositions[0]
+		status.WalApplyPosition = walPositions[0]
+		status.WalWritePosition = walPositions[0]
+	} else if len(walPositions) == 3 {
+		status = new(StandbyStatus)
+		status.WalFlushPosition = walPositions[0]
+		status.WalApplyPosition = walPositions[1]
+		status.WalWritePosition = walPositions[2]
+	} else {
+		err = errors.New(fmt.Sprintf("Invalid number of wal positions provided, need 1 or 3, got %d", len(walPositions)))
+		return
+	}
 	status.ClientTime = uint64((time.Now().UnixNano() - epochNano) / 1000)
 	return
 }
diff --git a/replication_test.go b/replication_test.go
index 411b449e..b86da951 100644
--- a/replication_test.go
+++ b/replication_test.go
@@ -150,9 +150,34 @@ func TestSimpleReplicationConnection(t *testing.T) {
 	// Before closing our connection, let's send a standby status to update our wal
 	// position, which should then be reflected if we fetch out our current wal position
 	// for the slot
-	replicationConn.SendStandbyStatus(pgx.NewStandbyStatus(maxWal))
+	status, err := pgx.NewStandbyStatus(maxWal)
+	if err != nil {
+		t.Errorf("Failed to create standby status %v", err)
+	}
+	replicationConn.SendStandbyStatus(status)
 	replicationConn.StopReplication()
 
+	// Let's push the boundary conditions of the standby status and ensure it errors correctly
+	status, err = pgx.NewStandbyStatus(0,1,2,3,4)
+	if err == nil {
+		t.Errorf("Expected error from new standby status, got %v",status)
+	}
+
+	// And if you provide 3 args, ensure the right fields are set
+	status, err = pgx.NewStandbyStatus(1,2,3)
+	if err != nil {
+		t.Errorf("Failed to create test status: %v", err)
+	}
+	if status.WalFlushPosition != 1 {
+		t.Errorf("Unexpected flush position %d", status.WalFlushPosition)
+	}
+	if status.WalApplyPosition != 2 {
+		t.Errorf("Unexpected apply position %d", status.WalApplyPosition)
+	}
+	if status.WalWritePosition != 3 {
+		t.Errorf("Unexpected write position %d", status.WalWritePosition)
+	}
+
 	err = replicationConn.Close()
 	if err != nil {
 		t.Fatalf("Replication connection close failed: %v", err)