From 417d61cf985d0e54434e4f66bc0a67d8d236a40f Mon Sep 17 00:00:00 2001
From: Masayuki Morita <minamijoyo@gmail.com>
Date: Thu, 21 Sep 2023 10:35:03 +0900
Subject: [PATCH] Wait until the function code has been successfully updated

Fixes #313

The current implementation does not wait LastUpdateStatus to be
Successful after UpdateFunctionCodeWithContext call. This might not
matter with small zip packages, but it does matter with large container
images.
---
 deploy.go | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/deploy.go b/deploy.go
index d9e99b0..6ada5a7 100644
--- a/deploy.go
+++ b/deploy.go
@@ -213,9 +213,11 @@ func (app *App) updateFunctionCode(ctx context.Context, in *lambda.UpdateFunctio
 		return nil, err
 	}
 
+	var res *lambda.FunctionConfiguration
 	retrier := retryPolicy.Start(ctx)
 	for retrier.Continue() {
-		res, err := app.lambda.UpdateFunctionCodeWithContext(ctx, in)
+		var err error
+		res, err = app.lambda.UpdateFunctionCodeWithContext(ctx, in)
 		if err != nil {
 			if aerr, ok := err.(awserr.Error); ok {
 				switch aerr.Code() {
@@ -226,10 +228,20 @@ func (app *App) updateFunctionCode(ctx context.Context, in *lambda.UpdateFunctio
 			}
 			return nil, errors.Wrap(err, "failed to update function code")
 		}
-		log.Println("[info] updated function code successfully")
-		return res, nil
+		log.Println("[info] update function code request was accepted")
+		break
 	}
-	return nil, errors.New("failed to update function code (max retries reached)")
+
+	if !retrier.Continue() {
+		return nil, errors.New("failed to update function code (max retries reached)")
+	}
+
+	if err := app.waitForLastUpdateStatusSuccessful(ctx, *in.FunctionName); err != nil {
+		return nil, err
+	}
+	log.Println("[info] updated function code successfully")
+
+	return res, nil
 }
 
 func (app *App) waitForLastUpdateStatusSuccessful(ctx context.Context, name string) error {