From 2b213771eba13c202f1794532d53642495abd71e Mon Sep 17 00:00:00 2001
From: Kishan B <kishancs46@gmail.com>
Date: Tue, 7 May 2024 17:09:34 +0530
Subject: [PATCH] Refactor to be more type sensitive

---
 main.go                       |  6 +++---
 pkg/internal/parser/parser.go |  2 +-
 pkg/tfimportgen.go            |  6 ++++--
 pkg/tfimportgen_test.go       | 23 ++++++++++++++++-------
 4 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/main.go b/main.go
index cb285ef..65ad304 100644
--- a/main.go
+++ b/main.go
@@ -42,11 +42,11 @@ terraform show -json | tf-import-gen aws_instance.example module.example
 terraform show -json | tf-import-gen
 `,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			address := ""
+			addresses := []string{""}
 			if len(os.Args) > 1 {
-				address = os.Args[1]
+				addresses = os.Args[1:]
 			}
-			imports, err := tfimportgen.GenerateImports(os.Stdin, address)
+			imports, err := tfimportgen.GenerateImports(os.Stdin, addresses)
 			if err != nil {
 				return err
 			}
diff --git a/pkg/internal/parser/parser.go b/pkg/internal/parser/parser.go
index e903258..95ef6a7 100644
--- a/pkg/internal/parser/parser.go
+++ b/pkg/internal/parser/parser.go
@@ -12,7 +12,7 @@ type TerraformResource struct {
 
 type TerraformResources []TerraformResource
 
-func (resources TerraformResources) FilterByAddresses(addresses ...string) TerraformResources {
+func (resources TerraformResources) FilterByAddresses(addresses []string) TerraformResources {
 	var filteredResources TerraformResources
 	for _, resource := range resources {
 		for _, address := range addresses {
diff --git a/pkg/tfimportgen.go b/pkg/tfimportgen.go
index 6b89bdc..703b7a1 100644
--- a/pkg/tfimportgen.go
+++ b/pkg/tfimportgen.go
@@ -5,13 +5,15 @@ import (
 	"io"
 )
 
-func GenerateImports(stateJsonReader io.Reader, address ...string) (TerraformImports, error) {
+func GenerateImports(stateJsonReader io.Reader, addresses []string) (TerraformImports, error) {
 	resources, err := parser.NewTerraformStateJsonParser(stateJsonReader).Parse()
 	if err != nil {
 		return nil, err
 	}
 
-	resources = resources.FilterByAddresses(address...)
+	if addresses != nil {
+		resources = resources.FilterByAddresses(addresses)
+	}
 
 	var imports TerraformImports
 	for _, resource := range resources {
diff --git a/pkg/tfimportgen_test.go b/pkg/tfimportgen_test.go
index 02dd432..86b53e1 100644
--- a/pkg/tfimportgen_test.go
+++ b/pkg/tfimportgen_test.go
@@ -1,8 +1,10 @@
 package tfimportgen_test
 
 import (
+	"bytes"
 	tfimportgen "github.com/kishaningithub/tf-import-gen/pkg"
 	"github.com/stretchr/testify/require"
+	"io"
 	"os"
 	"path/filepath"
 	"testing"
@@ -71,9 +73,16 @@ func Test_GenerateImports_ShouldGenerateImportsForAllResourcesWhenNoFiltersAreGi
 			t.Cleanup(func() {
 				_ = stateJsonFile.Close()
 			})
+			stateJsonFileBytes, err := io.ReadAll(stateJsonFile)
+			require.NoError(t, err)
+			stateJson := bytes.NewBuffer(stateJsonFileBytes)
 
-			actual, err := tfimportgen.GenerateImports(stateJsonFile, "")
+			actual, err := tfimportgen.GenerateImports(stateJson, []string{""})
+			require.NoError(t, err)
+			require.Equal(t, tt.expected, actual)
 
+			stateJson = bytes.NewBuffer(stateJsonFileBytes)
+			actual, err = tfimportgen.GenerateImports(stateJson, nil)
 			require.NoError(t, err)
 			require.Equal(t, tt.expected, actual)
 		})
@@ -83,12 +92,12 @@ func Test_GenerateImports_ShouldGenerateImportsForAllResourcesWhenNoFiltersAreGi
 func Test_GenerateImports_ShouldGenerateImportsForResourcesForGivenAddress(t *testing.T) {
 	tests := []struct {
 		name     string
-		address  string
+		address  []string
 		expected tfimportgen.TerraformImports
 	}{
 		{
 			name:    "filtering by module",
-			address: "module.test_mwaa",
+			address: []string{"module.test_mwaa"},
 			expected: tfimportgen.TerraformImports{
 				{
 					ResourceAddress: "module.test_mwaa.aws_iam_policy.test_mwaa_permissions",
@@ -104,7 +113,7 @@ func Test_GenerateImports_ShouldGenerateImportsForResourcesForGivenAddress(t *te
 		},
 		{
 			name:    "filtering by resource",
-			address: "aws_glue_catalog_database.test_db",
+			address: []string{"aws_glue_catalog_database.test_db"},
 			expected: tfimportgen.TerraformImports{
 				{
 					ResourceAddress: "aws_glue_catalog_database.test_db",
@@ -138,7 +147,7 @@ func Test_GenerateImports_ShouldGenerateHelpfulCommentForResourceThatCannotBeImp
 		_ = stateJsonFile.Close()
 	})
 
-	actual, err := tfimportgen.GenerateImports(stateJsonFile, "")
+	actual, err := tfimportgen.GenerateImports(stateJsonFile, []string{""})
 
 	require.NoError(t, err)
 	expectedImports := tfimportgen.TerraformImports{
@@ -156,7 +165,7 @@ func Test_GenerateImports_ShouldGenerateHelpfulCommentForResourceThatCannotBeImp
 	require.Equal(t, expectedImports, actual)
 }
 
-func Test_GenerateImports_ShouldGenerateImportsForResourcesForGivenAddresses(t *testing.T) {
+func Test_GenerateImports_ShouldGenerateImportsForResourcesForMultipleAddresses(t *testing.T) {
 	tests := []struct {
 		name     string
 		address  []string
@@ -193,7 +202,7 @@ func Test_GenerateImports_ShouldGenerateImportsForResourcesForGivenAddresses(t *
 				_ = stateJsonFile.Close()
 			})
 
-			actual, err := tfimportgen.GenerateImports(stateJsonFile, tt.address...)
+			actual, err := tfimportgen.GenerateImports(stateJsonFile, tt.address)
 
 			require.NoError(t, err)
 			require.Equal(t, tt.expected, actual)