From 3f5dfc957909e5e9fc9dc7e81fe8dc8a540a681b Mon Sep 17 00:00:00 2001
From: Brian Vander Schaaf <6099457+bvand@users.noreply.github.com>
Date: Fri, 9 Feb 2024 13:22:22 -0500
Subject: [PATCH] add 'elapsed' to applyProgress for deterministic waiting for
 heartbeat execution

---
 internal/command/views/hook_json.go      |  5 ++++
 internal/command/views/hook_json_test.go | 29 ++++++++++++++++++------
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/internal/command/views/hook_json.go b/internal/command/views/hook_json.go
index 9cfdbcee14b..4b785a67ecd 100644
--- a/internal/command/views/hook_json.go
+++ b/internal/command/views/hook_json.go
@@ -60,6 +60,8 @@ type applyProgress struct {
 	// heartbeatDone is used to allow tests to safely wait for the progress
 	// goroutine to finish
 	heartbeatDone chan struct{}
+
+	elapsed chan time.Duration
 }
 
 func (h *jsonHook) PreApply(addr addrs.AbsResourceInstance, gen states.Generation, action plans.Action, priorState, plannedNewState cty.Value) (tofu.HookAction, error) {
@@ -72,6 +74,7 @@ func (h *jsonHook) PreApply(addr addrs.AbsResourceInstance, gen states.Generatio
 		addr:          addr,
 		action:        action,
 		start:         h.timeNow().Round(time.Second),
+		elapsed:       make(chan time.Duration),
 		done:          make(chan struct{}),
 		heartbeatDone: make(chan struct{}),
 	}
@@ -87,6 +90,7 @@ func (h *jsonHook) PreApply(addr addrs.AbsResourceInstance, gen states.Generatio
 
 func (h *jsonHook) applyingHeartbeat(progress applyProgress) {
 	defer close(progress.heartbeatDone)
+	defer close(progress.elapsed)
 	for {
 		select {
 		case <-progress.done:
@@ -96,6 +100,7 @@ func (h *jsonHook) applyingHeartbeat(progress applyProgress) {
 
 		elapsed := h.timeNow().Round(time.Second).Sub(progress.start)
 		h.view.Hook(json.NewApplyProgress(progress.addr, progress.action, elapsed))
+		progress.elapsed <- elapsed
 	}
 }
 
diff --git a/internal/command/views/hook_json_test.go b/internal/command/views/hook_json_test.go
index ed766270cfe..44637b10635 100644
--- a/internal/command/views/hook_json_test.go
+++ b/internal/command/views/hook_json_test.go
@@ -9,6 +9,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
+
 	"github.com/opentofu/opentofu/internal/addrs"
 	"github.com/opentofu/opentofu/internal/plans"
 	"github.com/opentofu/opentofu/internal/states"
@@ -60,21 +62,25 @@ func TestJSONHook_create(t *testing.T) {
 	action, err = hook.PostProvisionInstanceStep(addr, "local-exec", nil)
 	testHookReturnValues(t, action, err)
 
-	// Travel 10s into the future, notify the progress goroutine, and sleep
-	// briefly to allow it to execute
+	elapsedChan := hook.applying[addr.String()].elapsed
+
+	// Travel 10s into the future, notify the progress goroutine, and wait
+	// for execution via 'elapsed' progress
 	nowMu.Lock()
 	now = now.Add(10 * time.Second)
 	after <- now
 	nowMu.Unlock()
-	time.Sleep(10 * time.Millisecond)
+	elapsed := <-elapsedChan
+	testDurationEqual(t, 10*time.Second, elapsed)
 
-	// Travel 10s into the future, notify the progress goroutine, and sleep
-	// briefly to allow it to execute
+	// Travel 10s into the future, notify the progress goroutine, and wait
+	// for execution via 'elapsed' progress
 	nowMu.Lock()
 	now = now.Add(10 * time.Second)
 	after <- now
 	nowMu.Unlock()
-	time.Sleep(10 * time.Millisecond)
+	elapsed = <-elapsedChan
+	testDurationEqual(t, 20*time.Second, elapsed)
 
 	// Travel 2s into the future. We have arrived!
 	nowMu.Lock()
@@ -83,11 +89,11 @@ func TestJSONHook_create(t *testing.T) {
 	action, err = hook.PostApply(addr, states.CurrentGen, plannedNewState, nil)
 	testHookReturnValues(t, action, err)
 
-
 	// Shut down the progress goroutine if still active
 	hook.applyingLock.Lock()
 	for key, progress := range hook.applying {
 		close(progress.done)
+		close(progress.elapsed)
 		<-progress.heartbeatDone
 		delete(hook.applying, key)
 	}
@@ -219,6 +225,7 @@ func TestJSONHook_errors(t *testing.T) {
 	hook.applyingLock.Lock()
 	for key, progress := range hook.applying {
 		close(progress.done)
+		close(progress.elapsed)
 		<-progress.heartbeatDone
 		delete(hook.applying, key)
 	}
@@ -339,3 +346,11 @@ func testHookReturnValues(t *testing.T, action tofu.HookAction, err error) {
 		t.Fatalf("Expected hook to continue, given: %#v", action)
 	}
 }
+
+func testDurationEqual(t *testing.T, wantedDuration time.Duration, gotDuration time.Duration) {
+	t.Helper()
+
+	if !cmp.Equal(wantedDuration, gotDuration) {
+		t.Errorf("unexpected time elapsed:%s\n", cmp.Diff(wantedDuration, gotDuration))
+	}
+}