From 191a5f1018f79010d2db986a7117dd91bd04c8e1 Mon Sep 17 00:00:00 2001
From: mrinalirao <mrinali.rao@hashicorp.com>
Date: Wed, 30 Nov 2022 14:31:10 +1100
Subject: [PATCH] add test for override

---
 internal/cloud/backend_taskStages_test.go | 61 +++++++++++++++++++++++
 internal/cloud/testing.go                 |  1 +
 internal/cloud/tfe_client_mock.go         | 32 ++++++++++++
 3 files changed, 94 insertions(+)

diff --git a/internal/cloud/backend_taskStages_test.go b/internal/cloud/backend_taskStages_test.go
index e52f6a5e70..e922b78331 100644
--- a/internal/cloud/backend_taskStages_test.go
+++ b/internal/cloud/backend_taskStages_test.go
@@ -3,6 +3,7 @@ package cloud
 import (
 	"context"
 	"errors"
+	"strings"
 	"testing"
 
 	"github.com/golang/mock/gomock"
@@ -205,3 +206,63 @@ func TestTaskStagesWithErrors(t *testing.T) {
 		t.Error("Expected to error but did not")
 	}
 }
+
+func TestTaskStageOverride(t *testing.T) {
+	b, bCleanup := testBackendWithName(t)
+	defer bCleanup()
+
+	integrationContext, writer := newMockIntegrationContext(b, t)
+
+	integrationContext.Op.UIOut = b.CLI
+
+	cases := map[string]struct {
+		taskStageID string
+		isError     bool
+		errMsg      string
+		input       *mockInput
+	}{
+		"override-pass": {
+			taskStageID: "ts-pass",
+			isError:     false,
+			input: testInput(t, map[string]string{
+				"→→ [bold]Override": "override",
+			}),
+			errMsg: "",
+		},
+		"override-fail": {
+			taskStageID: "ts-err",
+			isError:     true,
+			input: testInput(t, map[string]string{
+				"→→ [bold]Override": "override",
+			}),
+			errMsg: "",
+		},
+		"skip-override": {
+			taskStageID: "ts-err",
+			isError:     true,
+			errMsg:      "Failed to override: Apply discarded.",
+			input: testInput(t, map[string]string{
+				"→→ [bold]Override": "no",
+			}),
+		},
+	}
+	for _, c := range cases {
+		integrationContext.Op.UIIn = c.input
+		_, err := b.processStageOverrides(integrationContext, writer, c.taskStageID)
+		if c.isError {
+			if err == nil {
+				t.Fatalf("Expected to fail with some error")
+			}
+			if c.errMsg != "" {
+				if !strings.Contains(err.Error(), c.errMsg) {
+					t.Fatalf("Expected: %s, got: %s", c.errMsg, err.Error())
+				}
+			}
+
+		} else {
+			if err != nil {
+				t.Fatalf("Expected to pass, got err: %s", err)
+			}
+		}
+	}
+}
diff --git a/internal/cloud/testing.go b/internal/cloud/testing.go
index 9d2fc04452..10db6bdddf 100644
--- a/internal/cloud/testing.go
+++ b/internal/cloud/testing.go
@@ -209,6 +209,7 @@ func testBackend(t *testing.T, obj cty.Value) (*Cloud, func()) {
 	b.client.CostEstimates = mc.CostEstimates
 	b.client.Organizations = mc.Organizations
 	b.client.Plans = mc.Plans
+	b.client.TaskStages = mc.TaskStages
 	b.client.PolicySetOutcomes = mc.PolicySetOutcomes
 	b.client.PolicyChecks = mc.PolicyChecks
 	b.client.Runs = mc.Runs
diff --git a/internal/cloud/tfe_client_mock.go b/internal/cloud/tfe_client_mock.go
index 7d567833fb..815ce8c169 100644
--- a/internal/cloud/tfe_client_mock.go
+++ b/internal/cloud/tfe_client_mock.go
@@ -28,6 +28,7 @@ type MockClient struct {
 	Organizations         *MockOrganizations
 	Plans                 *MockPlans
 	PolicySetOutcomes     *MockPolicySetOutcomes
+	TaskStages            *MockTaskStages
 	PolicyChecks          *MockPolicyChecks
 	Runs                  *MockRuns
 	StateVersions         *MockStateVersions
@@ -43,6 +44,7 @@ func NewMockClient() *MockClient {
 	c.CostEstimates = newMockCostEstimates(c)
 	c.Organizations = newMockOrganizations(c)
 	c.Plans = newMockPlans(c)
+	c.TaskStages = newMockTaskStages(c)
 	c.PolicySetOutcomes = newMockPolicySetOutcomes(c)
 	c.PolicyChecks = newMockPolicyChecks(c)
 	c.Runs = newMockRuns(c)
@@ -547,6 +549,36 @@ func (m *MockPlans) ReadJSONOutput(ctx context.Context, planID string) ([]byte,
 	return []byte(planOutput), nil
 }
 
+type MockTaskStages struct {
+	client *MockClient
+}
+
+func newMockTaskStages(client *MockClient) *MockTaskStages {
+	return &MockTaskStages{
+		client: client,
+	}
+}
+
+func (m *MockTaskStages) Override(ctx context.Context, taskStageID string, options tfe.TaskStageOverrideOptions) (*tfe.TaskStage, error) {
+	switch taskStageID {
+	case "ts-err":
+		return nil, errors.New("test error")
+
+	default:
+		return nil, nil
+	}
+}
+
+func (m *MockTaskStages) Read(ctx context.Context, taskStageID string, options *tfe.TaskStageReadOptions) (*tfe.TaskStage, error) {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m *MockTaskStages) List(ctx context.Context, runID string, options *tfe.TaskStageListOptions) (*tfe.TaskStageList, error) {
+	//TODO implement me
+	panic("implement me")
+}
+
 type MockPolicySetOutcomes struct {
 	client *MockClient
 }