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

@@ -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.

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
+	var importId string
-	if importId == "" {
+	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)

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")
+	}
+}