From 7bba1d0c95dff08d6e115da3885306c2a9d6335a Mon Sep 17 00:00:00 2001
From: Chris Marchesi <chrism@vancluevertech.com>
Date: Wed, 6 Sep 2017 19:42:16 -0700
Subject: [PATCH] helper/resource: Add ImportStateIdFunc

Add an ImportStateIdFunc field to the ImportState testing functionality.
This will allow for more powerful generation of complex import state IDs
that can't be accomplished by ImportStateId or ImportStateIdPrefix
themselves.
---
 helper/resource/testing.go                   | 10 ++++
 helper/resource/testing_import_state.go      | 15 ++++-
 helper/resource/testing_import_state_test.go | 58 ++++++++++++++++++++
 3 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/helper/resource/testing.go b/helper/resource/testing.go
index 205547eff..8db38a77c 100644
--- a/helper/resource/testing.go
+++ b/helper/resource/testing.go
@@ -186,6 +186,10 @@ type TestCheckFunc func(*terraform.State) error
 // ImportStateCheckFunc is the check function for ImportState tests
 type ImportStateCheckFunc func([]*terraform.InstanceState) error
 
+// ImportStateIdFunc is an ID generation function to help with complex ID
+// generation for ImportState tests.
+type ImportStateIdFunc func(*terraform.State) (string, error)
+
 // TestCase is a single acceptance test case used to test the apply/destroy
 // lifecycle of a resource in a specific configuration.
 //
@@ -333,6 +337,12 @@ type TestStep struct {
 	// the unset ImportStateId field.
 	ImportStateIdPrefix string
 
+	// ImportStateIdFunc is a function that can be used to dynamically generate
+	// the ID for the ImportState tests. It is sent the state, which can be
+	// checked to derive the attributes necessary and generate the string in the
+	// desired format.
+	ImportStateIdFunc ImportStateIdFunc
+
 	// ImportStateCheck checks the results of ImportState. It should be
 	// used to verify that the resulting value of ImportState has the
 	// proper resources, IDs, and attributes.
diff --git a/helper/resource/testing_import_state.go b/helper/resource/testing_import_state.go
index 28ad10526..94fef3cfb 100644
--- a/helper/resource/testing_import_state.go
+++ b/helper/resource/testing_import_state.go
@@ -16,15 +16,24 @@ func testStepImportState(
 	state *terraform.State,
 	step TestStep) (*terraform.State, error) {
 	// Determine the ID to import
-	importId := step.ImportStateId
-	if importId == "" {
+	var importId string
+	switch {
+	case step.ImportStateIdFunc != nil:
+		var err error
+		importId, err = step.ImportStateIdFunc(state)
+		if err != nil {
+			return state, err
+		}
+	case step.ImportStateId != "":
+		importId = step.ImportStateId
+	default:
 		resource, err := testResource(step, state)
 		if err != nil {
 			return state, err
 		}
-
 		importId = resource.Primary.ID
 	}
+
 	importPrefix := step.ImportStateIdPrefix
 	if importPrefix != "" {
 		importId = fmt.Sprintf("%s%s", importPrefix, importId)
diff --git a/helper/resource/testing_import_state_test.go b/helper/resource/testing_import_state_test.go
index 8ef8323ba..c91355108 100644
--- a/helper/resource/testing_import_state_test.go
+++ b/helper/resource/testing_import_state_test.go
@@ -388,3 +388,61 @@ func TestTest_importStateVerifyFail(t *testing.T) {
 		t.Fatalf("test should fail")
 	}
 }
+
+func TestTest_importStateIdFunc(t *testing.T) {
+	mp := testProvider()
+	mp.ImportStateFn = func(
+		info *terraform.InstanceInfo, id string) ([]*terraform.InstanceState, error) {
+		if id != "foo:bar" {
+			return nil, fmt.Errorf("bad import ID: %s", id)
+		}
+
+		return []*terraform.InstanceState{
+			{
+				ID:        "foo",
+				Ephemeral: terraform.EphemeralState{Type: "test_instance"},
+			},
+		}, nil
+	}
+
+	mp.RefreshFn = func(
+		i *terraform.InstanceInfo,
+		s *terraform.InstanceState) (*terraform.InstanceState, error) {
+		return s, nil
+	}
+
+	checked := false
+	checkFn := func(s []*terraform.InstanceState) error {
+		checked = true
+
+		if s[0].ID != "foo" {
+			return fmt.Errorf("bad: %#v", s)
+		}
+
+		return nil
+	}
+
+	mt := new(mockT)
+	Test(mt, TestCase{
+		Providers: map[string]terraform.ResourceProvider{
+			"test": mp,
+		},
+
+		Steps: []TestStep{
+			TestStep{
+				Config:            testConfigStrProvider,
+				ResourceName:      "test_instance.foo",
+				ImportState:       true,
+				ImportStateIdFunc: func(*terraform.State) (string, error) { return "foo:bar", nil },
+				ImportStateCheck:  checkFn,
+			},
+		},
+	})
+
+	if mt.failed() {
+		t.Fatalf("test failed: %s", mt.failMessage())
+	}
+	if !checked {
+		t.Fatal("didn't call check")
+	}
+}