Create interface to allow retry unit test to be faster (#1014)

* Create interface to allow retry unit test to be faster

* Apply suggestions from code review

---------

Co-authored-by: Edward J <[email protected]>
Co-authored-by: Lio李歐 <[email protected]>

Author: 3 people

src/pkg/cli/client/byoc/baseclient.go

@@ -48,6 +48,8 @@ type ProjectBackend interface {
 }
 
 type ByocBaseClient struct {
+	client.RetryDelayer
+
 	PulumiStack             string
 	SetupDone               bool
 	ShouldDelegateSubdomain bool

src/pkg/cli/client/playground.go

@@ -12,6 +12,7 @@ import (
 
 type PlaygroundProvider struct {
 	GrpcClient
+	RetryDelayer
 }
 
 var _ Provider = (*PlaygroundProvider)(nil)

src/pkg/cli/client/provider.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"time"
 
+	"github.com/DefangLabs/defang/src/pkg"
 	"github.com/DefangLabs/defang/src/pkg/types"
 	defangv1 "github.com/DefangLabs/defang/src/protos/io/defang/v1"
 	composeTypes "github.com/compose-spec/compose-go/v2/types"
@@ -126,6 +128,7 @@ type Provider interface {
 	Delete(context.Context, *defangv1.DeleteRequest) (*defangv1.DeleteResponse, error)
 	DeleteConfig(context.Context, *defangv1.Secrets) error
 	Deploy(context.Context, *defangv1.DeployRequest) (*defangv1.DeployResponse, error)
+	DelayBeforeRetry(context.Context) error
 	Destroy(context.Context, *defangv1.DestroyRequest) (types.ETag, error)
 	Follow(context.Context, *defangv1.TailRequest) (ServerStream[defangv1.TailResponse], error)
 	GetService(context.Context, *defangv1.GetRequest) (*defangv1.ServiceInfo, error)
@@ -152,3 +155,14 @@ type Loader interface {
 	LoadProject(context.Context) (*composeTypes.Project, error)
 	LoadProjectName(context.Context) (string, error)
 }
+
+type RetryDelayer struct {
+	Delay time.Duration
+}
+
+func (r *RetryDelayer) DelayBeforeRetry(ctx context.Context) error {
+	if r.Delay == 0 {
+		r.Delay = 1 * time.Second // Minimum 1 second delay to be consistent with the old behavior
+	}
+	return pkg.SleepWithContext(ctx, r.Delay)
+}

src/pkg/cli/subscribe.go

@@ -3,7 +3,6 @@ package cli
 import (
 	"context"
 	"errors"
-	"time"
 
 	"github.com/DefangLabs/defang/src/pkg"
 	"github.com/DefangLabs/defang/src/pkg/cli/client"
@@ -48,7 +47,9 @@ func WaitServiceState(
 		if !serverStream.Receive() {
 			// Reconnect on Error: internal: stream error: stream ID 5; INTERNAL_ERROR; received from peer
 			if isTransientError(serverStream.Err()) {
-				pkg.SleepWithContext(ctx, 1*time.Second)
+				if err := provider.DelayBeforeRetry(ctx); err != nil {
+					return err
+				}
 				serverStream, err = provider.Subscribe(ctx, &subscribeRequest)
 				if err != nil {
 					return err

src/pkg/cli/subscribe_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/DefangLabs/defang/src/pkg"
 	"github.com/DefangLabs/defang/src/pkg/cli/client"
@@ -267,6 +268,7 @@ func (m *MockSubscribeServerStreamForReconnectTest) Err() error {
 type mockSubscribeProviderForReconnectTest struct {
 	client.MockProvider
 	stream *MockSubscribeServerStreamForReconnectTest
+	client.RetryDelayer
 }
 
 func (m *mockSubscribeProviderForReconnectTest) Subscribe(
@@ -316,7 +318,7 @@ func TestWaitServiceStateStreamReceive(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			ctx := context.Background()
-			provider := &mockSubscribeProviderForReconnectTest{stream: tt.stream}
+			provider := &mockSubscribeProviderForReconnectTest{stream: tt.stream, RetryDelayer: client.RetryDelayer{Delay: 1 * time.Millisecond}}
 			err := WaitServiceState(
 				ctx, provider,
 				defangv1.ServiceState_DEPLOYMENT_COMPLETED,