From 9493db0fbdde86c2266f2719a8eb5a124305a21e Mon Sep 17 00:00:00 2001
From: Adam Mckaig <adam.mckaig@gmail.com>
Date: Thu, 18 Aug 2016 22:15:55 -0400
Subject: [PATCH] Check .terraform/terraform.tfstate, for remote state (#41)

* Move input path stuff to GetInputPath
* Look for .terraform/terraform.tfstate
* Update README
---
 README.md     |  7 ++---
 input.go      | 35 ++++++++++++++++++++++++
 input_test.go | 75 +++++++++++++++++++++++++++++++++++++++++++++++++++
 main.go       | 21 ++++-----------
 4 files changed, 119 insertions(+), 19 deletions(-)
 create mode 100644 input.go
 create mode 100644 input_test.go

diff --git a/README.md b/README.md
index 95279aa..736a245 100644
--- a/README.md
+++ b/README.md
@@ -24,8 +24,8 @@ ready to go.
 
 ## Usage
 
-If your Terraform state file is named `terraform.tfstate` (the default), `cd` to
-it and run:
+If you are using [remote state][rs] (or if your state file happens to be named
+`terraform.tfstate`), `cd` to it and run:
 
 	ansible-playbook --inventory-file=/path/to/terraform-inventory deploy/playbook.yml
 
@@ -139,4 +139,5 @@ MIT.
 
 [ansible]: https://www.ansible.com
 [tf]:      https://www.terraform.io
-[rel]:    https://github.com/adammck/terraform-inventory/releases
+[rel]:     https://github.com/adammck/terraform-inventory/releases
+[rs]:      https://www.terraform.io/docs/state/remote/index.html
diff --git a/input.go b/input.go
new file mode 100644
index 0000000..ca7b8bb
--- /dev/null
+++ b/input.go
@@ -0,0 +1,35 @@
+package main
+
+import (
+	"github.com/adammck/venv"
+	"github.com/blang/vfs"
+)
+
+func GetInputPath(fs vfs.Filesystem, env venv.Env) string {
+
+	var fn string
+
+	fn = env.Getenv("TF_STATE")
+	if fn != "" {
+		return fn
+	}
+
+	fn = env.Getenv("TI_TFSTATE")
+	if fn != "" {
+		return fn
+	}
+
+	fn = "terraform.tfstate"
+	_, err := fs.Stat(fn)
+	if err == nil {
+		return fn
+	}
+
+	fn = ".terraform/terraform.tfstate"
+	_, err = fs.Stat(fn)
+	if err == nil {
+		return fn
+	}
+
+	return ""
+}
diff --git a/input_test.go b/input_test.go
new file mode 100644
index 0000000..78b23f1
--- /dev/null
+++ b/input_test.go
@@ -0,0 +1,75 @@
+package main
+
+import (
+	"io"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/adammck/venv"
+	"github.com/blang/vfs"
+	"github.com/blang/vfs/memfs"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetInputPath(t *testing.T) {
+	assert.Equal(t, "", GetInputPath(memfs.Create(), venv.Mock()))
+	assert.Equal(t, "aaa", GetInputPath(memfs.Create(), envWith(map[string]string{"TF_STATE": "aaa"})))
+	assert.Equal(t, "bbb", GetInputPath(memfs.Create(), envWith(map[string]string{"TI_TFSTATE": "bbb"})))
+	assert.Equal(t, "terraform.tfstate", GetInputPath(fsWithFiles([]string{"terraform.tfstate"}), venv.Mock()))
+	assert.Equal(t, ".terraform/terraform.tfstate", GetInputPath(fsWithFiles([]string{".terraform/terraform.tfstate"}), venv.Mock()))
+}
+
+func envWith(env map[string]string) venv.Env {
+	e := venv.Mock()
+
+	for k, v := range env {
+		e.Setenv(k, v)
+	}
+
+	return e
+}
+
+func fsWithFiles(filenames []string) vfs.Filesystem {
+	fs := memfs.Create()
+	var err error
+
+	for _, fn := range filenames {
+
+		path := filepath.Dir(fn)
+		if path != "" {
+			err = vfs.MkdirAll(fs, path, 0700)
+			if err != nil {
+				panic(err)
+			}
+		}
+
+		err = touchFile(fs, fn)
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	return fs
+}
+
+// TODO: Upgrade this later with file contents.
+func touchFile(fs vfs.Filesystem, filename string) error {
+	return writeFile(fs, filename, []byte{}, 0600)
+}
+
+// port of ioutil.Writefile for vfs
+func writeFile(fs vfs.Filesystem, filename string, data []byte, perm os.FileMode) error {
+	f, err := fs.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
+	if err != nil {
+		return err
+	}
+	n, err := f.Write(data)
+	if err == nil && n < len(data) {
+		err = io.ErrShortWrite
+	}
+	if err1 := f.Close(); err == nil {
+		err = err1
+	}
+	return err
+}
diff --git a/main.go b/main.go
index 7d1edff..d534f62 100644
--- a/main.go
+++ b/main.go
@@ -3,6 +3,8 @@ package main
 import (
 	"flag"
 	"fmt"
+	"github.com/adammck/venv"
+	"github.com/blang/vfs"
 	"os"
 	"path/filepath"
 )
@@ -21,23 +23,10 @@ func main() {
 		return
 	}
 
-	// not given on the command line? try ENV.
 	if file == "" {
-		file = os.Getenv("TF_STATE")
-	}
-
-	// also try the old ENV name.
-	if file == "" {
-		file = os.Getenv("TI_TFSTATE")
-	}
-
-	// check for a file named terraform.tfstate in the pwd
-	if file == "" {
-		fn := "terraform.tfstate"
-		_, err := os.Stat(fn)
-		if err == nil {
-			file = fn
-		}
+		fs := vfs.OS()
+		env := venv.OS()
+		file = GetInputPath(fs, env)
 	}
 
 	if file == "" {