diff --git a/tfwrite/file_test.go b/tfwrite/file_test.go
index be6b97c..e1636b4 100644
--- a/tfwrite/file_test.go
+++ b/tfwrite/file_test.go
@@ -4,15 +4,68 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 )
 
-func TestFileFindResourcesByType(t *testing.T) {
+func TestParseBlock(t *testing.T) {
 	cases := []struct {
-		desc         string
-		src          string
-		resourceType string
-		want         string
-		ok           bool
+		desc string
+		src  string
+		want string
+		ok   bool
+	}{
+		{
+			desc: "resource",
+			src: `
+resource "foo_test" "example" {}
+`,
+			want: "resource",
+			ok:   true,
+		},
+		{
+			desc: "data",
+			src: `
+data "foo_test" "example" {}
+`,
+			want: "data",
+			ok:   true,
+		},
+		{
+			desc: "provider",
+			src: `
+provider "foo" {}
+`,
+			want: "provider",
+			ok:   true,
+		},
+		{
+			desc: "unknown",
+			src: `
+b "foo" {}
+`,
+			want: "b",
+			ok:   true,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.desc, func(t *testing.T) {
+			f := parseTestFile(t, tc.src)
+			b := parseBlock(findFirstTestBlock(t, f).Raw())
+			got := b.Type()
+			if got != tc.want {
+				t.Errorf("got = %s, but want = %s", got, tc.want)
+			}
+		})
+	}
+}
+
+func TestFileBlocks(t *testing.T) {
+	cases := []struct {
+		desc string
+		src  string
+		want []Block
+		ok   bool
 	}{
 		{
 			desc: "simple",
@@ -22,7 +75,66 @@ resource "foo_test" "example2" {}
 resource "foo_bar" "example1" {}
 data "foo_test" "example1" {}
 `,
-			resourceType: "foo_test",
+			want: []Block{
+				NewEmptyResource("foo_test", "example1"),
+				NewEmptyResource("foo_test", "example2"),
+				NewEmptyResource("foo_bar", "example1"),
+				NewEmptyDataSource("foo_test", "example1"),
+			},
+			ok: true,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.desc, func(t *testing.T) {
+			f := parseTestFile(t, tc.src)
+			got := f.Blocks()
+			opts := cmpopts.IgnoreUnexported(Resource{}, DataSource{})
+			if diff := cmp.Diff(got, tc.want, opts); diff != "" {
+				t.Errorf("got:\n%s\nwant:\n%s\ndiff:\n%s", got, tc.want, diff)
+			}
+		})
+	}
+}
+
+func TestFileFindBlocksByType(t *testing.T) {
+	cases := []struct {
+		desc       string
+		src        string
+		blockType  string
+		schemaType string
+		want       string
+		ok         bool
+	}{
+		{
+			desc: "resource",
+			src: `
+resource "foo_test" "example1" {}
+resource "foo_test" "example2" {}
+resource "foo_bar" "example1" {}
+data "foo_test" "example1" {}
+`,
+			blockType:  "resource",
+			schemaType: "",
+			want: `
+resource "foo_test" "example1" {}
+
+resource "foo_test" "example2" {}
+
+resource "foo_bar" "example1" {}
+`,
+			ok: true,
+		},
+		{
+			desc: "resource with schemaType",
+			src: `
+resource "foo_test" "example1" {}
+resource "foo_test" "example2" {}
+resource "foo_bar" "example1" {}
+data "foo_test" "example1" {}
+`,
+			blockType:  "resource",
+			schemaType: "foo_test",
 			want: `
 resource "foo_test" "example1" {}
 
@@ -35,57 +147,60 @@ resource "foo_test" "example2" {}
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			resources := f.FindResourcesByType(tc.resourceType)
+			blocks := f.FindBlocksByType(tc.blockType, tc.schemaType)
 
 			newFile := NewEmptyFile()
-			for _, r := range resources {
+			for _, r := range blocks {
 				newFile.AppendBlock(r)
 			}
 
 			got := printTestFile(t, newFile)
 			if diff := cmp.Diff(got, tc.want); diff != "" {
-				t.Fatalf("got:\n%s\nwant:\n%s\ndiff:\n%s", got, tc.want, diff)
+				t.Errorf("got:\n%s\nwant:\n%s\ndiff:\n%s", got, tc.want, diff)
 			}
 		})
 	}
 }
 
-func TestFileFindProvidersByType(t *testing.T) {
+func TestFileFindResourcesByType(t *testing.T) {
 	cases := []struct {
-		desc         string
-		src          string
-		providerType string
-		want         string
-		ok           bool
+		desc       string
+		src        string
+		schemaType string
+		want       string
+		ok         bool
 	}{
 		{
-			desc: "simple",
+			desc: "resource",
 			src: `
-provider "aws" {
-  region = "ap-northeast-1"
-  alias  = "ap_northeast_1"
-}
-
-provider "aws" {
-  region = "us-east-1"
-  alias  = "us_east_1"
-}
+resource "foo_test" "example1" {}
+resource "foo_test" "example2" {}
+resource "foo_bar" "example1" {}
+data "foo_test" "example1" {}
+`,
+			schemaType: "",
+			want: `
+resource "foo_test" "example1" {}
 
-provider "google" {}
+resource "foo_test" "example2" {}
 
-resource "aws_s3_bucket" "example" {}
+resource "foo_bar" "example1" {}
+`,
+			ok: true,
+		},
+		{
+			desc: "resource with schemaType",
+			src: `
+resource "foo_test" "example1" {}
+resource "foo_test" "example2" {}
+resource "foo_bar" "example1" {}
+data "foo_test" "example1" {}
 `,
-			providerType: "aws",
+			schemaType: "foo_test",
 			want: `
-provider "aws" {
-  region = "ap-northeast-1"
-  alias  = "ap_northeast_1"
-}
+resource "foo_test" "example1" {}
 
-provider "aws" {
-  region = "us-east-1"
-  alias  = "us_east_1"
-}
+resource "foo_test" "example2" {}
 `,
 			ok: true,
 		},
@@ -94,16 +209,16 @@ provider "aws" {
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			f := parseTestFile(t, tc.src)
-			providers := f.FindProvidersByType(tc.providerType)
+			blocks := f.FindResourcesByType(tc.schemaType)
 
 			newFile := NewEmptyFile()
-			for _, p := range providers {
-				newFile.AppendBlock(p)
+			for _, r := range blocks {
+				newFile.AppendBlock(r)
 			}
 
 			got := printTestFile(t, newFile)
 			if diff := cmp.Diff(got, tc.want); diff != "" {
-				t.Fatalf("got:\n%s\nwant:\n%s\ndiff:\n%s", got, tc.want, diff)
+				t.Errorf("got:\n%s\nwant:\n%s\ndiff:\n%s", got, tc.want, diff)
 			}
 		})
 	}