From 4e25f014d8ba9bd88df3317ec51ce3fa783203d0 Mon Sep 17 00:00:00 2001
From: Megan Wolf <97549300+meganwolf0@users.noreply.github.com>
Date: Thu, 3 Oct 2024 10:47:01 -0400
Subject: [PATCH] fix(console): reset compdef when editing (#701)

---
 src/internal/testhelpers/testhelpers.go           | 11 +++++++----
 .../assessment_results/assessment-results_test.go |  2 +-
 src/internal/tui/component/component_test.go      | 15 +++++++++++----
 src/internal/tui/model_test.go                    |  2 +-
 4 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/src/internal/testhelpers/testhelpers.go b/src/internal/testhelpers/testhelpers.go
index e119ff75..c9aa9c5b 100644
--- a/src/internal/testhelpers/testhelpers.go
+++ b/src/internal/testhelpers/testhelpers.go
@@ -37,14 +37,14 @@ func CreateTempFile(t *testing.T, ext string) *os.File {
 }
 
 // RunTestModelView runs a test model view with a given model and messages, impelements a retry loop if final model is nil
-func RunTestModelView(t *testing.T, m tea.Model, msgs []tea.Msg, timeout time.Duration, maxRetries, height, width int) error {
+func RunTestModelView(t *testing.T, m tea.Model, reset func() tea.Model, msgs []tea.Msg, timeout time.Duration, maxRetries, height, width int) error {
 
-	testModelView := func(t *testing.T) (bool, error) {
+	testModelView := func(t *testing.T, try int) (bool, error) {
 		tm := teatest.NewTestModel(t, m, teatest.WithInitialTermSize(width, height))
 
 		for _, msg := range msgs {
 			tm.Send(msg)
-			time.Sleep(time.Millisecond * 50)
+			time.Sleep(time.Millisecond * time.Duration(50*try))
 		}
 
 		if err := tm.Quit(); err != nil {
@@ -63,8 +63,11 @@ func RunTestModelView(t *testing.T, m tea.Model, msgs []tea.Msg, timeout time.Du
 	}
 
 	for i := 0; i < maxRetries; i++ {
-		retry, err := testModelView(t)
+		retry, err := testModelView(t, i+1)
 		if retry {
+			if reset != nil {
+				m = reset()
+			}
 			continue
 		}
 		if err != nil {
diff --git a/src/internal/tui/assessment_results/assessment-results_test.go b/src/internal/tui/assessment_results/assessment-results_test.go
index e57240a0..5dc56898 100644
--- a/src/internal/tui/assessment_results/assessment-results_test.go
+++ b/src/internal/tui/assessment_results/assessment-results_test.go
@@ -33,7 +33,7 @@ func TestAssessmentResultsBasicView(t *testing.T) {
 
 	msgs := []tea.Msg{}
 
-	err := testhelpers.RunTestModelView(t, model, msgs, timeout, maxRetries, height, width)
+	err := testhelpers.RunTestModelView(t, model, nil, msgs, timeout, maxRetries, height, width)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/src/internal/tui/component/component_test.go b/src/internal/tui/component/component_test.go
index ed263a30..fa85ea3c 100644
--- a/src/internal/tui/component/component_test.go
+++ b/src/internal/tui/component/component_test.go
@@ -39,7 +39,7 @@ func TestComponentDefinitionBasicView(t *testing.T) {
 
 	msgs := []tea.Msg{}
 
-	err := testhelpers.RunTestModelView(t, model, msgs, timeout, maxRetries, height, width)
+	err := testhelpers.RunTestModelView(t, model, nil, msgs, timeout, maxRetries, height, width)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -61,7 +61,7 @@ func TestComponentDefinitionComponentSwitch(t *testing.T) {
 		tea.KeyMsg{Type: tea.KeyEnter}, // Open control
 	}
 
-	err := testhelpers.RunTestModelView(t, model, msgs, timeout, maxRetries, height, width)
+	err := testhelpers.RunTestModelView(t, model, nil, msgs, timeout, maxRetries, height, width)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -81,7 +81,7 @@ func TestComponentControlSelect(t *testing.T) {
 		tea.KeyMsg{Type: tea.KeyEnter}, // Open control
 	}
 
-	err := testhelpers.RunTestModelView(t, model, msgs, timeout, maxRetries, height, width)
+	err := testhelpers.RunTestModelView(t, model, nil, msgs, timeout, maxRetries, height, width)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -105,7 +105,14 @@ func TestEditViewComponentDefinitionModel(t *testing.T) {
 		tea.KeyMsg{Type: tea.KeyEnter},                                    // Confirm edit
 	}
 
-	err := testhelpers.RunTestModelView(t, model, msgs, timeout, maxRetries, height, width)
+	reset := func() tea.Model {
+		resetOscalModel := testhelpers.OscalFromPath(t, validCompDefValidations)
+		resetModel := component.NewComponentDefinitionModel(resetOscalModel.ComponentDefinition)
+		resetModel.Open(height, width)
+		return resetModel
+	}
+
+	err := testhelpers.RunTestModelView(t, model, reset, msgs, timeout, maxRetries, height, width)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/src/internal/tui/model_test.go b/src/internal/tui/model_test.go
index e3a3cfd6..c366beb2 100644
--- a/src/internal/tui/model_test.go
+++ b/src/internal/tui/model_test.go
@@ -34,7 +34,7 @@ func TestNewOSCALModel(t *testing.T) {
 
 	msgs := []tea.Msg{}
 
-	err := testhelpers.RunTestModelView(t, model, msgs, timeout, maxRetries, height, width)
+	err := testhelpers.RunTestModelView(t, model, nil, msgs, timeout, maxRetries, height, width)
 	if err != nil {
 		t.Fatal(err)
 	}