From bcebcb55a2756321d521d889f5e8bf4451c16268 Mon Sep 17 00:00:00 2001
From: Jan De Dobbeleer <jan.de.dobbeleer@gmail.com>
Date: Fri, 1 Dec 2023 20:48:40 +0100
Subject: [PATCH] feat(terraform): add all contextual files

---
 src/segments/terraform.go      | 51 ++++++++++++++++++++++----------
 src/segments/terraform_test.go | 54 +++++++++++++++++++++-------------
 2 files changed, 69 insertions(+), 36 deletions(-)

diff --git a/src/segments/terraform.go b/src/segments/terraform.go
index d4a9ad2b1ddd..6c1671a7b140 100644
--- a/src/segments/terraform.go
+++ b/src/segments/terraform.go
@@ -39,53 +39,73 @@ type TerraformBlock struct {
 
 func (tf *Terraform) Enabled() bool {
 	cmd := "terraform"
-	terraformFolder := filepath.Join(tf.env.Pwd(), ".terraform")
 	fetchVersion := tf.props.GetBool(properties.FetchVersion, false)
-	if fetchVersion {
-		// known version files
-		files := []string{"versions.tf", "main.tf", "terraform.tfstate"}
-		var hasFiles bool
-		for _, file := range files {
-			if tf.env.HasFiles(file) {
-				hasFiles = true
-				break
-			}
-		}
-		fetchVersion = hasFiles
-	}
 
-	inContext := tf.env.HasFolder(terraformFolder) || fetchVersion
-	if !tf.env.HasCommand(cmd) || !inContext {
+	if !tf.env.HasCommand(cmd) || !tf.inContext(fetchVersion) {
 		return false
 	}
+
 	tf.WorkspaceName, _ = tf.env.RunCommand(cmd, "workspace", "show")
 	if !fetchVersion {
 		return true
 	}
+
 	if err := tf.setVersionFromTfFiles(); err == nil {
 		return true
 	}
+
 	tf.setVersionFromTfStateFile()
 	return true
 }
 
+func (tf *Terraform) inContext(fetchVersion bool) bool {
+	terraformFolder := filepath.Join(tf.env.Pwd(), ".terraform")
+
+	if tf.env.HasFolder(terraformFolder) {
+		return true
+	}
+
+	files := []string{".tf", ".tfplan", ".tfstate"}
+	for _, file := range files {
+		if tf.env.HasFiles(file) {
+			return true
+		}
+	}
+
+	if !fetchVersion {
+		return false
+	}
+
+	versionFiles := []string{"versions.tf", "main.tf", "terraform.tfstate"}
+	for _, file := range versionFiles {
+		if tf.env.HasFiles(file) {
+			return true
+		}
+	}
+
+	return false
+}
+
 func (tf *Terraform) setVersionFromTfFiles() error {
 	files := []string{"versions.tf", "main.tf"}
 	for _, file := range files {
 		if !tf.env.HasFiles(file) {
 			continue
 		}
+
 		parser := hclparse.NewParser()
 		content := tf.env.FileContent(file)
 		hclFile, diags := parser.ParseHCL([]byte(content), file)
 		if diags != nil {
 			continue
 		}
+
 		var config TerraFormConfig
 		diags = gohcl.DecodeBody(hclFile.Body, nil, &config)
 		if diags != nil {
 			continue
 		}
+
 		tf.TerraformBlock = *config.Terraform
 		return nil
 	}
@@ -97,6 +117,7 @@ func (tf *Terraform) setVersionFromTfStateFile() {
 	if !tf.env.HasFiles(file) {
 		return
 	}
+
 	content := tf.env.FileContent(file)
 	_ = json.Unmarshal([]byte(content), &tf.TerraformBlock)
 }
diff --git a/src/segments/terraform_test.go b/src/segments/terraform_test.go
index 3a2c2486cf7f..ec0629c746f7 100644
--- a/src/segments/terraform_test.go
+++ b/src/segments/terraform_test.go
@@ -12,16 +12,17 @@ import (
 
 func TestTerraform(t *testing.T) {
 	cases := []struct {
-		Case            string
-		Template        string
-		HasTfCommand    bool
-		HasTfFolder     bool
-		HasTfFiles      bool
-		HasTfStateFile  bool
-		FetchVersion    bool
-		WorkspaceName   string
-		ExpectedString  string
-		ExpectedEnabled bool
+		Case              string
+		Template          string
+		HasTfCommand      bool
+		HasTfFolder       bool
+		HasTfFiles        bool
+		HasTfVersionFiles bool
+		HasTfStateFile    bool
+		FetchVersion      bool
+		WorkspaceName     string
+		ExpectedString    string
+		ExpectedEnabled   bool
 	}{
 		{
 			Case:            "default workspace",
@@ -51,24 +52,32 @@ func TestTerraform(t *testing.T) {
 			FetchVersion:   true,
 		},
 		{
-			Case:            "files",
-			ExpectedString:  ">= 1.0.10",
+			Case:              "files",
+			ExpectedString:    ">= 1.0.10",
+			ExpectedEnabled:   true,
+			WorkspaceName:     "default",
+			Template:          "{{ .Version }}",
+			HasTfVersionFiles: true,
+			HasTfCommand:      true,
+			FetchVersion:      true,
+		},
+		{
+			Case:            "version files",
+			ExpectedString:  "0.12.24",
 			ExpectedEnabled: true,
 			WorkspaceName:   "default",
 			Template:        "{{ .Version }}",
-			HasTfFiles:      true,
+			HasTfStateFile:  true,
 			HasTfCommand:    true,
 			FetchVersion:    true,
 		},
 		{
-			Case:            "files",
-			ExpectedString:  "0.12.24",
+			Case:            "context files",
+			ExpectedString:  "default",
 			ExpectedEnabled: true,
 			WorkspaceName:   "default",
-			Template:        "{{ .Version }}",
-			HasTfStateFile:  true,
+			HasTfFiles:      true,
 			HasTfCommand:    true,
-			FetchVersion:    true,
 		},
 	}
 
@@ -77,12 +86,15 @@ func TestTerraform(t *testing.T) {
 
 		env.On("HasCommand", "terraform").Return(tc.HasTfCommand)
 		env.On("HasFolder", ".terraform").Return(tc.HasTfFolder)
+		env.On("HasFiles", ".tf").Return(tc.HasTfFiles)
+		env.On("HasFiles", ".tfplan").Return(tc.HasTfFiles)
+		env.On("HasFiles", ".tfstate").Return(tc.HasTfFiles)
 		env.On("Pwd").Return("")
 		env.On("RunCommand", "terraform", []string{"workspace", "show"}).Return(tc.WorkspaceName, nil)
-		env.On("HasFiles", "versions.tf").Return(tc.HasTfFiles)
-		env.On("HasFiles", "main.tf").Return(tc.HasTfFiles)
+		env.On("HasFiles", "versions.tf").Return(tc.HasTfVersionFiles)
+		env.On("HasFiles", "main.tf").Return(tc.HasTfVersionFiles)
 		env.On("HasFiles", "terraform.tfstate").Return(tc.HasTfStateFile)
-		if tc.HasTfFiles {
+		if tc.HasTfVersionFiles {
 			content, _ := os.ReadFile("../test/versions.tf")
 			env.On("FileContent", "versions.tf").Return(string(content))
 		}