Fix test

planetscale · Jan 21, 2025 · 975eff6 · 975eff6
1 parent 7079623
commit 975eff6
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 6 deletions.
diff --git a/internal/cmd/keyspace/rollout_status.go b/internal/cmd/keyspace/rollout_status.go
@@ -2,6 +2,7 @@ package keyspace
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/planetscale/cli/internal/cmdutil"
 	"github.com/planetscale/cli/internal/printer"
@@ -71,10 +72,17 @@ type ShardRollout struct {
 }
 
 func toShardRollout(sr ps.ShardRollout) *ShardRollout {
+	var startedAt, finishedAt *time.Time
+	if !sr.LastRolloutStartedAt.IsZero() {
+		startedAt = &sr.LastRolloutStartedAt
+	}
+	if !sr.LastRolloutFinishedAt.IsZero() {
+		finishedAt = &sr.LastRolloutFinishedAt
+	}
 	return &ShardRollout{
 		Name:                  sr.Name,
 		State:                 cmdutil.SnakeToSentenceCase(sr.State),
-		LastRolloutStartedAt:  printer.GetMillisecondsIfExists(&sr.LastRolloutStartedAt),
-		LastRolloutFinishedAt: printer.GetMillisecondsIfExists(&sr.LastRolloutFinishedAt),
+		LastRolloutStartedAt:  printer.GetMillisecondsIfExists(startedAt),
+		LastRolloutFinishedAt: printer.GetMillisecondsIfExists(finishedAt),
 	}
 }
diff --git a/internal/cmd/keyspace/rollout_status_test.go b/internal/cmd/keyspace/rollout_status_test.go
@@ -32,23 +32,38 @@ func TestKeyspace_RolloutStatusCmd(t *testing.T) {
 
 	kr := &ps.KeyspaceRollout{
 		Name:  "sharded",
-		State: "complete",
+		State: "resizing",
 		Shards: []ps.ShardRollout{
 			{
 				Name:                  "-80",
-				State:                 "complete",
+				State:                 "resizing",
 				LastRolloutFinishedAt: time.Time{},
 				LastRolloutStartedAt:  ts,
 			},
 			{
 				Name:                  "80-",
-				State:                 "complete",
+				State:                 "resizing",
 				LastRolloutFinishedAt: time.Time{},
 				LastRolloutStartedAt:  ts,
 			},
 		},
 	}
 
+	expected := []ShardRollout{
+		{
+			Name:                  "-80",
+			State:                 cmdutil.SnakeToSentenceCase("resizing"),
+			LastRolloutFinishedAt: nil,
+			LastRolloutStartedAt:  printer.GetMillisecondsIfExists(&ts),
+		},
+		{
+			Name:                  "80-",
+			State:                 cmdutil.SnakeToSentenceCase("resizing"),
+			LastRolloutFinishedAt: nil,
+			LastRolloutStartedAt:  printer.GetMillisecondsIfExists(&ts),
+		},
+	}
+
 	svc := &mock.KeyspacesService{
 		RolloutStatusFn: func(ctx context.Context, req *ps.KeyspaceRolloutStatusRequest) (*ps.KeyspaceRollout, error) {
 			c.Assert(req.Database, qt.Equals, db)
@@ -77,5 +92,5 @@ func TestKeyspace_RolloutStatusCmd(t *testing.T) {
 	err := cmd.Execute()
 	c.Assert(err, qt.IsNil)
 	c.Assert(svc.RolloutStatusFnInvoked, qt.IsTrue)
-	c.Assert(buf.String(), qt.JSONEquals, kr.Shards)
+	c.Assert(buf.String(), qt.JSONEquals, expected)
 }