Address PR review comments

internal/planner/planner.go

@@ -258,10 +258,8 @@ func getScaleDeployments(
 				if d.Spec.Replicas != nil && *d.Spec.Replicas != replicas {
 					scaleDeployments[version.Deployment] = uint32(replicas)
 				}
-			} else {
-				if d.Spec.Replicas != nil && *d.Spec.Replicas != 0 {
-					scaleDeployments[version.Deployment] = 0
-				}
+			} else if d.Spec.Replicas != nil && *d.Spec.Replicas != 0 {
+				scaleDeployments[version.Deployment] = 0
 			}
 		case temporaliov1alpha1.VersionStatusRamping, temporaliov1alpha1.VersionStatusCurrent:
 			// Scale up these deployments

internal/testhelpers/example_test.go

@@ -0,0 +1,71 @@
+package testhelpers_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	temporaliov1alpha1 "github.com/temporalio/temporal-worker-controller/api/v1alpha1"
+	"github.com/temporalio/temporal-worker-controller/internal/testhelpers"
+)
+
+// ExampleTestCase demonstrates how to create and use a TestCase for integration testing.
+// This example shows the basic structure and usage patterns for TestCase.
+func ExampleTestCase() {
+	// Create a test case using the builder pattern
+	testCaseBuilder := testhelpers.NewTestCase().
+		WithInput(
+			testhelpers.NewTemporalWorkerDeploymentBuilder().
+				WithName("example-worker").
+				WithNamespace("default").
+				WithManualStrategy().
+				WithTargetTemplate("v1"),
+		).
+		WithExpectedStatus(
+			testhelpers.NewStatusBuilder().
+				WithTargetVersion("v1", temporaliov1alpha1.VersionStatusInactive, -1, true, false),
+		).
+		WithWaitTime(5 * time.Second).
+		WithSetupFunction(func(t *testing.T, ctx context.Context, tc testhelpers.TestCase, env testhelpers.TestEnv) {
+			// Custom setup logic can go here
+			fmt.Println("Setting up test environment")
+		})
+
+	testCase := testCaseBuilder.BuildWithValues("example-worker", "default", "default")
+
+	// Access the TestCase fields
+	twd := testCase.GetTWD()
+	fmt.Printf("Testing deployment: %s/%s\n", twd.Namespace, twd.Name)
+	fmt.Printf("Strategy: %s\n", twd.Spec.RolloutStrategy.Strategy)
+
+	// Output:
+	// Testing deployment: default/example-worker
+	// Strategy: Manual
+}
+
+// ExampleTestCase_withExistingDeployments demonstrates how to create a TestCase
+// that starts with existing deprecated deployments.
+func ExampleTestCase_withExistingDeployments() {
+	testCaseBuilder := testhelpers.NewTestCase().
+		WithInput(
+			testhelpers.NewTemporalWorkerDeploymentBuilder().
+				WithName("worker-with-history").
+				WithNamespace("test-ns").
+				WithAllAtOnceStrategy().
+				WithTargetTemplate("v2"),
+		).
+		WithExistingDeployments(testhelpers.DeploymentInfo{}).
+		WithExpectedStatus(
+			testhelpers.NewStatusBuilder().
+				WithTargetVersion("v2", temporaliov1alpha1.VersionStatusInactive, -1, true, false),
+		)
+
+	testCase := testCaseBuilder.BuildWithValues("worker-with-history", "test-ns", "default")
+
+	twd := testCase.GetTWD()
+	fmt.Printf("Worker: %s/%s\n", twd.Namespace, twd.Name)
+
+	// Output:
+	// Worker: test-ns/worker-with-history
+}

internal/testhelpers/make.go

@@ -13,7 +13,7 @@ import (
 )
 
 const (
-	TaskQueueEnvKey = "TEMPORAL_TASK_QUEUE"
+	taskQueueEnvKey = "TEMPORAL_TASK_QUEUE"
 )
 
 func MakeTWD(
@@ -64,7 +64,7 @@ func MakeTWD(
 // MakePodSpec creates a pod spec with the given containers, labels, and task queue
 func MakePodSpec(containers []corev1.Container, labels map[string]string, taskQueue string) corev1.PodTemplateSpec {
 	for i := range containers {
-		containers[i].Env = append(containers[i].Env, corev1.EnvVar{Name: TaskQueueEnvKey, Value: taskQueue})
+		containers[i].Env = append(containers[i].Env, corev1.EnvVar{Name: taskQueueEnvKey, Value: taskQueue})
 	}
 
 	return corev1.PodTemplateSpec{
@@ -84,19 +84,19 @@ func MakePodSpecWithImage(imageName string) corev1.PodTemplateSpec {
 		"")
 }
 
-// SetTaskQueue sets or replaces the env var TaskQueueEnvKey with the given string in all containers
+// SetTaskQueue sets or replaces the env var taskQueueEnvKey with the given string in all containers
 func SetTaskQueue(podSpec corev1.PodTemplateSpec, taskQueue string) corev1.PodTemplateSpec {
 	for i, c := range podSpec.Spec.Containers {
 		found := false
 		for j, e := range c.Env {
-			if e.Name == TaskQueueEnvKey {
+			if e.Name == taskQueueEnvKey {
 				found = true
 				podSpec.Spec.Containers[i].Env[j].Value = taskQueue
 			}
 		}
 		if !found {
 			podSpec.Spec.Containers[i].Env = append(podSpec.Spec.Containers[i].Env, corev1.EnvVar{
-				Name:  TaskQueueEnvKey,
+				Name:  taskQueueEnvKey,
 				Value: taskQueue,
 			})
 		}

internal/testhelpers/test_builder.go

@@ -209,25 +209,37 @@ func (sb *StatusBuilder) Build() *temporaliov1alpha1.TemporalWorkerDeploymentSta
 	return ret
 }
 
+// TestCase represents a single test scenario for integration testing of TemporalWorkerDeployment controllers.
+// It encapsulates the input state, expected outputs, and any additional setup required for the test.
 type TestCase struct {
-	// If starting from a particular state, specify that in input.Status
+	// twd is the TemporalWorkerDeployment resource to test with.
+	// If starting from a particular state, specify that in input.Status.
 	twd *temporaliov1alpha1.TemporalWorkerDeployment
+
+	// existingDeploymentReplicas specifies the number of replicas for each deprecated build.
 	// TemporalWorkerDeploymentStatus only tracks the names of the Deployments for deprecated
 	// versions, so for test scenarios that start with existing deprecated version Deployments,
 	// specify the number of replicas for each deprecated build here.
 	existingDeploymentReplicas map[string]int32
+
+	// existingDeploymentImages specifies the images for each deprecated build.
 	// TemporalWorkerDeploymentStatus only tracks the build ids of the Deployments for deprecated
 	// versions, not their images so for test scenarios that start with existing deprecated version Deployments,
 	// specify the images for each deprecated build here.
 	existingDeploymentImages map[string]string
-	expectedStatus           *temporaliov1alpha1.TemporalWorkerDeploymentStatus
-	// validate that deployments have correct # of replicas. TODO(carlydf): validate replica count for more than just the deprecated versions
+
+	// expectedStatus is the expected TemporalWorkerDeploymentStatus after the test completes.
+	expectedStatus *temporaliov1alpha1.TemporalWorkerDeploymentStatus
+
+	// expectedDeploymentReplicas validates that deployments have correct number of replicas.
+	// TODO(carlydf): validate replica count for more than just the deprecated versions
 	expectedDeploymentReplicas map[string]int32
-	// Time to delay before checking expected status
+
+	// waitTime is the duration to delay before checking expected status.
 	waitTime *time.Duration
 
-	// Arbitrary function called at the end of setting up the environment specified by input.Status.
-	// Can be used for additional state creation / destruction
+	// setupFunc is an arbitrary function called at the end of setting up the environment specified by input.Status.
+	// It can be used for additional state creation or destruction.
 	setupFunc func(t *testing.T, ctx context.Context, tc TestCase, env TestEnv)
 }
 

internal/testhelpers/workers.go

@@ -37,7 +37,7 @@ func newVersionedWorker(ctx context.Context, podTemplateSpec corev1.PodTemplateS
 	if err != nil {
 		return nil, nil, err
 	}
-	temporalTaskQueue, err := getEnv(podTemplateSpec, TaskQueueEnvKey)
+	temporalTaskQueue, err := getEnv(podTemplateSpec, taskQueueEnvKey)
 	if err != nil {
 		return nil, nil, err
 	}

internal/tests/README.md

@@ -75,28 +75,7 @@ If that becomes an issue, it can be changed, but I think using one `temporaltest
 functional test and sharing it across test cases will make the tests run a lot faster as we add more 
 test cases.
 
-Each test case should pass in:
-```go
-type TestCase struct {
-	// If starting from a particular state, specify that in input.Status
-	twd *temporaliov1alpha1.TemporalWorkerDeployment
-	// TemporalWorkerDeploymentStatus only tracks the names of the Deployments for deprecated
-	// versions, so for test scenarios that start with existing deprecated version Deployments,
-	// specify the number of replicas for each deprecated build here.
-	existingDeploymentReplicas map[string]int32
-	// TemporalWorkerDeploymentStatus only tracks the build ids of the Deployments for deprecated
-	// versions, not their images so for test scenarios that start with existing deprecated version Deployments,
-	// specify the images for each deprecated build here.
-	existingDeploymentImages map[string]string
-	expectedStatus           *temporaliov1alpha1.TemporalWorkerDeploymentStatus
-	// Time to delay before checking expected status
-	waitTime *time.Duration
-
-	// Arbitrary function called at the end of setting up the environment specified by input.Status.
-	// Can be used for additional state creation / destruction
-	setupFunc func(t *testing.T, ctx context.Context, tc TestCase, env TestEnv)
-}
-```
+Each test case should use the `testhelpers.TestCase` struct. See the godoc for `testhelpers.TestCase` for detailed field descriptions and the examples in `testhelpers/example_test.go` for usage patterns.
 
 Each test function should follow the same pattern:
 1. Set up test environment (including preliminary state)