Refactor block-related logic

tfwrite/data_source_test.go

@@ -24,9 +24,9 @@ data "foo_test" "example" {}
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			r := findFirstTestDataSource(t, f)
+			b := NewDataSource(findFirstTestBlock(t, f).Raw())
 
-			got := r.Type()
+			got := b.Type()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}
@@ -54,8 +54,8 @@ data "foo_test" "example" {}
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			r := findFirstTestDataSource(t, f)
-			got := r.SchemaType()
+			b := NewDataSource(findFirstTestBlock(t, f).Raw())
+			got := b.SchemaType()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}
@@ -83,8 +83,8 @@ data "foo_test" "example" {}
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			r := findFirstTestDataSource(t, f)
-			got := r.Name()
+			b := NewDataSource(findFirstTestBlock(t, f).Raw())
+			got := b.Name()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}
@@ -133,8 +133,8 @@ data "foo_test" "example" {
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			r := findFirstTestDataSource(t, f)
-			got := r.ReferableName()
+			b := NewDataSource(findFirstTestBlock(t, f).Raw())
+			got := b.ReferableName()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}

tfwrite/file.go

@@ -23,89 +23,89 @@ func (f *File) Raw() *hclwrite.File {
 	return f.raw
 }
 
+// findBlocksByType returns all matching blocks from the body that have the
+// given blockType or returns an empty list if not found.
+func (f *File) findBlocksByType(blockType string) []*block {
+	var blocks []*block
+
+	for _, block := range f.Raw().Body().Blocks() {
+		if block.Type() == blockType {
+			blocks = append(blocks, newBlock(block))
+		}
+	}
+
+	return blocks
+}
+
 // FindResourcesByType returns all matching resources from the body that have the
 // given resourceType or returns an empty list if not found.
-func (f *File) FindResourcesByType(resourceType string) []*Resource {
+func (f *File) FindResourcesByType(schemaType string) []*Resource {
 	var matched []*Resource
 
-	for _, block := range f.raw.Body().Blocks() {
-		if block.Type() != "resource" {
+	for _, block := range f.findBlocksByType("resource") {
+		b := NewResource(block.Raw())
+		if b.SchemaType() != schemaType {
 			continue
 		}
 
-		labels := block.Labels()
-		if len(labels) == 2 && labels[0] != resourceType {
-			continue
-		}
-
-		resource := NewResource(block)
-		matched = append(matched, resource)
+		matched = append(matched, b)
 	}
 
 	return matched
 }
 
 // FindDataSourcesByType returns all matching data sources from the body that have the
 // given dataSourceType or returns an empty list if not found.
-func (f *File) FindDataSourcesByType(dataSourceType string) []*DataSource {
+func (f *File) FindDataSourcesByType(schemaType string) []*DataSource {
 	var matched []*DataSource
 
-	for _, block := range f.raw.Body().Blocks() {
-		if block.Type() != "data" {
-			continue
-		}
-
-		labels := block.Labels()
-		if len(labels) == 2 && labels[0] != dataSourceType {
+	for _, block := range f.findBlocksByType("data") {
+		b := NewDataSource(block.Raw())
+		if b.SchemaType() != schemaType {
 			continue
 		}
 
-		dataSource := NewDataSource(block)
-		matched = append(matched, dataSource)
+		matched = append(matched, b)
 	}
 
 	return matched
 }
 
 // FindProvidersByType returns all matching providers from the body that have the
 // given providerType or returns an empty list if not found.
-func (f *File) FindProvidersByType(providerType string) []*Provider {
+func (f *File) FindProvidersByType(schemaType string) []*Provider {
 	var matched []*Provider
 
-	for _, block := range f.raw.Body().Blocks() {
-		if block.Type() != "provider" {
+	for _, block := range f.findBlocksByType("provider") {
+		b := NewProvider(block.Raw())
+		if b.SchemaType() != schemaType {
 			continue
 		}
 
-		labels := block.Labels()
-		if len(labels) == 1 && labels[0] != providerType {
-			continue
-		}
-
-		provider := NewProvider(block)
-		matched = append(matched, provider)
+		matched = append(matched, b)
 	}
 
 	return matched
 }
 
-// AppendResource appends a given resource to the file.
-func (f *File) AppendResource(resource *Resource) {
+// appendBlock appends a given block to the file.
+func (f *File) appendBlock(block Block) {
 	body := f.raw.Body()
 	body.AppendNewline()
-	body.AppendBlock(resource.raw)
+	body.AppendBlock(block.Raw())
+}
+
+// AppendResource appends a given resource to the file.
+func (f *File) AppendResource(block *Resource) {
+	f.appendBlock(block)
 }
 
 // AppendDataSource appends a given data source to the file.
-func (f *File) AppendDataSource(dataSource *DataSource) {
-	body := f.raw.Body()
-	body.AppendNewline()
-	body.AppendBlock(dataSource.raw)
+func (f *File) AppendDataSource(block *DataSource) {
+	f.appendBlock(block)
 }
 
 // AppendProvider appends a given provider to the file.
-func (f *File) AppendProvider(provider *Provider) {
-	body := f.raw.Body()
-	body.AppendNewline()
-	body.AppendBlock(provider.raw)
+func (f *File) AppendProvider(block *Provider) {
+	f.appendBlock(block)
 }

tfwrite/helper_test.go

@@ -44,36 +44,3 @@ func findTestBlocks(t *testing.T, f *File) []*block {
 	}
 	return blocks
 }
-
-// findFirstTestResource is a test helper for find the first resource.
-func findFirstTestResource(t *testing.T, f *File) *Resource {
-	t.Helper()
-	for _, block := range f.raw.Body().Blocks() {
-		if block.Type() == "resource" {
-			return NewResource(block)
-		}
-	}
-	return nil
-}
-
-// findFirstTestDataSource is a test helper for find the first data source.
-func findFirstTestDataSource(t *testing.T, f *File) *DataSource {
-	t.Helper()
-	for _, block := range f.raw.Body().Blocks() {
-		if block.Type() == "data" {
-			return NewDataSource(block)
-		}
-	}
-	return nil
-}
-
-// findFirstTestProvider is a test helper for find the first provider.
-func findFirstTestProvider(t *testing.T, f *File) *Provider {
-	t.Helper()
-	for _, block := range f.raw.Body().Blocks() {
-		if block.Type() == "provider" {
-			return NewProvider(block)
-		}
-	}
-	return nil
-}

tfwrite/provider_test.go

@@ -24,9 +24,9 @@ provider "foo" {}
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			p := findFirstTestProvider(t, f)
+			b := NewProvider(findFirstTestBlock(t, f).Raw())
 
-			got := p.Type()
+			got := b.Type()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}
@@ -54,8 +54,8 @@ provider "foo" {}
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			p := findFirstTestProvider(t, f)
-			got := p.SchemaType()
+			b := NewProvider(findFirstTestBlock(t, f).Raw())
+			got := b.SchemaType()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}

tfwrite/resource_test.go

@@ -30,9 +30,9 @@ resource "foo_test" "example" {
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			r := findFirstTestResource(t, f)
+			b := NewResource(findFirstTestBlock(t, f).Raw())
 
-			got := r.Type()
+			got := b.Type()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}
@@ -60,8 +60,8 @@ resource "aws_s3_bucket" "example" {}
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			r := findFirstTestResource(t, f)
-			got := r.SchemaType()
+			b := NewResource(findFirstTestBlock(t, f).Raw())
+			got := b.SchemaType()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}
@@ -89,8 +89,8 @@ resource "aws_s3_bucket" "example" {}
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			r := findFirstTestResource(t, f)
-			got := r.Name()
+			b := NewResource(findFirstTestBlock(t, f).Raw())
+			got := b.Name()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}
@@ -139,8 +139,8 @@ resource "aws_s3_bucket" "example" {
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			r := findFirstTestResource(t, f)
-			got := r.ReferableName()
+			b := NewResource(findFirstTestBlock(t, f).Raw())
+			got := b.ReferableName()
 			if got != tc.want {
 				t.Errorf("got = %s, but want = %s", got, tc.want)
 			}