From 6fb0eaf97a27ce00610e614bbfb56f5854d13bee Mon Sep 17 00:00:00 2001
From: Martin Atkins <mart@degeneration.co.uk>
Date: Sat, 18 Nov 2017 19:27:38 -0800
Subject: [PATCH] testharness: generate multiple child contexts for counted
 resource

When we encounter a describe block for a resource address that refers to
a resource with count != 1, a child context is created for each of the
instances of that resource, causing the associated testers to be
automatically applied to each instance in turn.
---
 testharness/context.go         | 20 ++++++++---
 testharness/context_setters.go | 62 +++++++++++++++++++++++++++++++---
 2 files changed, 72 insertions(+), 10 deletions(-)

diff --git a/testharness/context.go b/testharness/context.go
index 4b6d6bf947..a7f193be66 100644
--- a/testharness/context.go
+++ b/testharness/context.go
@@ -3,6 +3,7 @@ package testharness
 import (
 	"fmt"
 
+	"github.com/hashicorp/terraform/terraform"
 	lua "github.com/yuin/gopher-lua"
 	"github.com/zclconf/go-cty/cty"
 )
@@ -50,12 +51,14 @@ func (ctx *Context) WithName(name string) *Context {
 // WithNameSuffix returns a new context that has the given string appended to
 // the name of the receiving context.
 func (ctx *Context) WithNameSuffix(suffix string) *Context {
+	return ctx.WithName(ctx.nameWithSuffix(suffix))
+}
+
+func (ctx *Context) nameWithSuffix(suffix string) string {
 	if ctx.name == "" {
-		return ctx.WithName(suffix)
+		return suffix
 	}
-	retVal := *ctx
-	retVal.name = fmt.Sprintf("%s %s", ctx.name, suffix)
-	return &retVal
+	return fmt.Sprintf("%s %s", ctx.name, suffix)
 }
 
 // HasResource returns true if there is a resource object associated with
@@ -72,8 +75,9 @@ func (ctx *Context) Resource() cty.Value {
 
 // WithResource returns a new context which has the given resource object
 // associated.
-func (ctx *Context) WithResource(obj cty.Value) *Context {
+func (ctx *Context) WithResource(addr *terraform.ResourceAddress, obj cty.Value) *Context {
 	retVal := *ctx
+	retVal.name = ctx.nameWithSuffix(addr.String())
 	retVal.resource = obj
 	return &retVal
 }
@@ -152,3 +156,9 @@ func (ctx *Context) EachObject() cty.Value {
 	}
 	return cty.ObjectVal(ctx.each)
 }
+
+func (ctx *Context) withNewLuaThread() *Context {
+	retVal := *ctx
+	retVal.lstate, _ = ctx.lstate.NewThread()
+	return &retVal
+}
diff --git a/testharness/context_setters.go b/testharness/context_setters.go
index b07e0aca88..fc87e3cafa 100644
--- a/testharness/context_setters.go
+++ b/testharness/context_setters.go
@@ -1,8 +1,11 @@
 package testharness
 
 import (
+	"fmt"
+
 	"github.com/hashicorp/terraform/terraform"
 	"github.com/hashicorp/terraform/tfdiags"
+	"github.com/zclconf/go-cty/cty"
 )
 
 // A contextSetter is something passed to the first argument of a "describe"
@@ -41,9 +44,58 @@ type resourceContextSetter struct {
 }
 
 func (s *resourceContextSetter) AppendContexts(parent *Context, subject *Subject, ctxs []*Context) ([]*Context, tfdiags.Diagnostics) {
-	// TODO: Set the resource object too
-	// TODO: If the resource address refers to a resource block with multiple
-	// instances (e.g. "count" is set) then generate one context for each
-	// of the instances matched.
-	return append(ctxs, parent.WithNameSuffix(s.Addr.String())), nil
+	var diags tfdiags.Diagnostics
+
+	// FIXME: Check if a resource with the given address is defined _at all_
+	// and return an error if not. When we do this, we must handle the special
+	// case where the resource _is_ defined but has count = 0, in which case
+	// that is not an error but rather we just generate no child contexts
+	// at all.
+
+	filter := &terraform.StateFilter{
+		State: subject.state,
+	}
+	// The StateFilter interface is weird: it expects ResourceAddress _strings_
+	// which it parses itself, rather than letting the caller do its own
+	// validation. Since we already parsed and validated our resource address,
+	// we'll need to stringify it here and let this function re-parse it. :/
+	results, err := filter.Filter(s.Addr.String())
+	if err != nil {
+		// The only possible error is an invalid address, which should never
+		// happen because we're passing in a pre-validated address here.
+		// Thus we won't go to any unusual effort to make this a user-friendly
+		// diagnostic.
+		diags = diags.Append(err)
+		return ctxs, diags
+	}
+
+	for _, result := range results {
+		// Again, for some reason the StateFilter interface deals in strings
+		// rather than ResourceAddress objects, so once again we end up
+		// redundantly round-tripping through a string. :(
+		addr, err := terraform.ParseResourceAddress(result.Address)
+		if err != nil {
+			// Should never happen because this address was just handed
+			// to us by Terraform core
+			panic(fmt.Errorf("invalid resource address generated by Terraform core: %s", err))
+		}
+
+		var inst *terraform.InstanceState
+		switch tr := result.Value.(type) {
+		case *terraform.InstanceState:
+			inst = tr
+		default:
+			// should never happen, but if it does we'll ignore it since
+			// it's presumably some new type of thing in state.
+			continue
+		}
+
+		// TODO: convert inst into a cty.Value representing the instance,
+		// which we can then place in the context for downstream test
+		// code to use.
+		_ = inst
+		ctxs = append(ctxs, parent.WithResource(addr, cty.DynamicVal))
+	}
+
+	return ctxs, diags
 }