Fix tests; add WaitForImageRegionStatus

linode · lgarber-akamai · Jun 28, 2024 · Jun 13, 2024 · Jun 25, 2024 · Jun 25, 2024
commit 02417029c84239b2b1150e873f48a633ac581c5e
@@ -20,11 +20,11 @@ const (
 	ImageStatusAvailable     ImageStatus = "available"
 )
 
-// ImageRegionStatus represents the status of an Image's replication.
+// ImageRegionStatus represents the status of an Image's replica.
 type ImageRegionStatus string
 
 // ImageRegionStatus options start with ImageRegionStatus and
-// include all Image replication statuses
+// include all Image replica statuses
 const (
 	ImageRegionStatusAvailable          ImageRegionStatus = "available"
 	ImageRegionStatusCreating           ImageRegionStatus = "creating"

@@ -3,7 +3,6 @@ package integration
 import (
 	"bytes"
 	"context"
-	"fmt"
 	"slices"
 	"testing"
 
@@ -113,7 +112,7 @@ func TestImage_CreateUpload(t *testing.T) {
 
 	image, uploadURL, err := client.CreateImageUpload(context.Background(), ImageCreateUploadOptions{
 		// TODO: Uncomment before merging Images Gen. 2 support to main
-		//Region:      getRegionsWithCaps(t, client, []string{"Metadata"})[0],
+		// Region:      getRegionsWithCaps(t, client, []string{"Metadata"})[0],
 		Region: "us-east",
 
 		Label:       "linodego-image-create-upload",
@@ -144,7 +143,7 @@ func TestImage_CloudInit(t *testing.T) {
 		t, "fixtures/TestImage_CloudInit", true,
 		func(client *Client, options *InstanceCreateOptions) {
 			// TODO: Uncomment before merging Images Gen. 2 support to main
-			//options.Region = getRegionsWithCaps(t, client, []string{"Metadata"})[0]
+			// options.Region = getRegionsWithCaps(t, client, []string{"Metadata"})[0]
 			options.Region = "us-east"
 		})
 	if err != nil {
@@ -220,24 +219,20 @@ func TestImage_Replicate(t *testing.T) {
 
 	image, err = client.WaitForImageStatus(context.Background(), image.ID, ImageStatusAvailable, 240)
 	require.NoError(t, err)
-	//
-	//require.Equal(t, testRegion, image.Regions[0].Region)
-	//require.NotZero(t, image.Regions[0].Status)
+
+	replicaRegions := availableRegions[1:]
 
 	image, err = client.ReplicateImage(context.Background(), image.ID, ImageReplicateOptions{
-		Regions: availableRegions[1:],
+		Regions: replicaRegions,
 	})
 	require.NoError(t, err)
 
-	// Wait for the image to start replicating
-	image, err = client.WaitForImageStatus(context.Background(), image.ID, ImageStatusReplicating, 240)
-	require.NoError(t, err)
-
-	// Wait for the replication process to complete
-	image, err = client.WaitForImageStatus(context.Background(), image.ID, ImageStatusAvailable, 240)
-	require.NoError(t, err)
-
-	fmt.Println(image)
+	// Wait for all region replications to finish
+	for _, region := range replicaRegions {
+		image, err = client.WaitForImageRegionStatus(
+			context.Background(), image.ID, region, ImageRegionStatusAvailable,
+		)
+	}
 
 	require.Len(t, image.Regions, len(availableRegions))
 	for _, region := range image.Regions {

@@ -3,13 +3,14 @@ package integration
 import (
 	"context"
 	"fmt"
-	"github.com/linode/linodego"
 	"io"
 	"log"
 	"net/http"
 	"os"
 	"testing"
 	"time"
+
+	"github.com/linode/linodego"
 )
 
 var (

@@ -25,15 +25,15 @@ func TestImage_Replicate(t *testing.T) {
 		Regions: []linodego.ImageRegion{
 			{
 				Region: "us-iad",
-				Status: linodego.ImageStatusAvailable,
+				Status: linodego.ImageRegionStatusAvailable,
 			},
 			{
 				Region: "us-mia",
-				Status: linodego.ImageStatusReplicating,
+				Status: linodego.ImageRegionStatusReplicating,
 			},
 			{
 				Region: "us-ord",
-				Status: linodego.ImageStatusPendingReplication,
+				Status: linodego.ImageRegionStatusPendingReplication,
 			},
 		},
 	}
@@ -51,11 +51,11 @@ func TestImage_Replicate(t *testing.T) {
 	require.Equal(t, "test", image.Label)
 
 	require.EqualValues(t, "us-iad", image.Regions[0].Region)
-	require.EqualValues(t, linodego.ImageStatusAvailable, image.Regions[0].Status)
+	require.EqualValues(t, linodego.ImageRegionStatusAvailable, image.Regions[0].Status)
 
 	require.EqualValues(t, "us-mia", image.Regions[1].Region)
-	require.EqualValues(t, linodego.ImageStatusReplicating, image.Regions[1].Status)
+	require.EqualValues(t, linodego.ImageRegionStatusReplicating, image.Regions[1].Status)
 
 	require.EqualValues(t, "us-ord", image.Regions[2].Region)
-	require.EqualValues(t, linodego.ImageStatusPendingReplication, image.Regions[2].Status)
+	require.EqualValues(t, linodego.ImageRegionStatusPendingReplication, image.Regions[2].Status)
 }
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"net/http"
+	"slices"
 	"strconv"
 	"time"
 
@@ -442,6 +443,40 @@ func (client Client) WaitForImageStatus(ctx context.Context, imageID string, sta
 	}
 }
 
+// WaitForImageRegionStatus waits for an Image's replica to reach the desired state
+// before returning.
+func (client Client) WaitForImageRegionStatus(ctx context.Context, imageID, region string, status ImageRegionStatus) (*Image, error) {
+	ticker := time.NewTicker(client.pollInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			image, err := client.GetImage(ctx, imageID)
+			if err != nil {
+				return image, err
+			}
+
+			replicaIdx := slices.IndexFunc(
+				image.Regions,
+				func(r ImageRegion) bool {
+					return r.Region == region
+				},
+			)
+
+			// If no replica was found or the status doesn't match, try again
+			if replicaIdx < 0 || image.Regions[replicaIdx].Status != status {
+				continue
+			}
+
+			return image, nil
+
+		case <-ctx.Done():
+			return nil, fmt.Errorf("failed to wait for Image %s status %s: %w", imageID, status, ctx.Err())
+		}
+	}
+}
+
 // WaitForMySQLDatabaseBackup waits for the backup with the given label to be available.
 func (client Client) WaitForMySQLDatabaseBackup(ctx context.Context, dbID int, label string, timeoutSeconds int) (*MySQLDatabaseBackup, error) {
 	ctx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second)