Allow count in data sources

builtin/providers/test/data_source.go

@@ -11,11 +11,21 @@ func testDataSource() *schema.Resource {
 		Read: testDataSourceRead,
 
 		Schema: map[string]*schema.Schema{
-			"list": &schema.Schema{
+			"list": {
 				Type:     schema.TypeList,
 				Computed: true,
 				Elem:     &schema.Schema{Type: schema.TypeString},
 			},
+
+			"input": {
+				Type:     schema.TypeString,
+				Optional: true,
+			},
+
+			"output": {
+				Type:     schema.TypeString,
+				Computed: true,
+			},
 		},
 	}
 }
@@ -24,5 +34,11 @@ func testDataSourceRead(d *schema.ResourceData, meta interface{}) error {
 	d.SetId(time.Now().UTC().String())
 	d.Set("list", []interface{}{"one", "two", "three"})
 
+	if input, hasInput := d.GetOk("input"); hasInput {
+		d.Set("output", input)
+	} else {
+		d.Set("output", "some output")
+	}
+
 	return nil
 }

builtin/providers/test/data_source_test.go

@@ -0,0 +1,57 @@
+package test
+
+import (
+	"errors"
+	"strings"
+	"testing"
+
+	"github.com/hashicorp/terraform/helper/resource"
+	"github.com/hashicorp/terraform/terraform"
+)
+
+func TestDataSource_dataSourceCount(t *testing.T) {
+	resource.UnitTest(t, resource.TestCase{
+		Providers: testAccProviders,
+		CheckDestroy: func(s *terraform.State) error {
+			return nil
+		},
+		Steps: []resource.TestStep{
+			{
+				Config: strings.TrimSpace(`
+data "test_data_source" "test" {
+  count = 3
+  input = "count-${count.index}"
+}
+
+resource "test_resource" "foo" {
+  required     = "yep"
+  required_map = {
+    key = "value"
+  }
+
+  list = ["${data.test_data_source.test.*.output}"]
+}
+				`),
+				Check: func(s *terraform.State) error {
+					res, hasRes := s.RootModule().Resources["test_resource.foo"]
+					if !hasRes {
+						return errors.New("No test_resource.foo in state")
+					}
+					if res.Primary.Attributes["list.#"] != "3" {
+						return errors.New("Wrong list.#, expected 3")
+					}
+					if res.Primary.Attributes["list.0"] != "count-0" {
+						return errors.New("Wrong list.0, expected count-0")
+					}
+					if res.Primary.Attributes["list.1"] != "count-1" {
+						return errors.New("Wrong list.0, expected count-1")
+					}
+					if res.Primary.Attributes["list.2"] != "count-2" {
+						return errors.New("Wrong list.0, expected count-2")
+					}
+					return nil
+				},
+			},
+		},
+	})
+}

builtin/providers/test/resource.go

@@ -97,6 +97,13 @@ func testResource() *schema.Resource {
 				Type:     schema.TypeMap,
 				Computed: true,
 			},
+			"list": {
+				Type:     schema.TypeList,
+				Optional: true,
+				Elem: &schema.Schema{
+					Type: schema.TypeString,
+				},
+			},
 			"list_of_map": {
 				Type:     schema.TypeList,
 				Optional: true,

config/config_test.go

@@ -552,3 +552,20 @@ func testConfig(t *testing.T, name string) *Config {
 
 	return c
 }
+
+func TestConfigDataCount(t *testing.T) {
+	c := testConfig(t, "data-count")
+	actual, err := c.Resources[0].Count()
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if actual != 5 {
+		t.Fatalf("bad: %#v", actual)
+	}
+
+	// we need to make sure "count" has been removed from the RawConfig, since
+	// it's not a real key and won't validate.
+	if _, ok := c.Resources[0].RawConfig.Raw["count"]; ok {
+		t.Fatal("count key still exists in RawConfig")
+	}
+}

config/loader_hcl.go

@@ -457,6 +457,7 @@ func loadDataResourcesHcl(list *ast.ObjectList) ([]*Resource, error) {
 		// Remove the fields we handle specially
 		delete(config, "depends_on")
 		delete(config, "provider")
+		delete(config, "count")
 
 		rawConfig, err := NewRawConfig(config)
 		if err != nil {

config/test-fixtures/data-count/main.tf

@@ -0,0 +1,3 @@
+data "foo" "bar" {
+    count = 5
+}

terraform/context_plan_test.go

@@ -911,6 +911,37 @@ func TestContext2Plan_computedDataResource(t *testing.T) {
 	}
 }
 
+func TestContext2Plan_computedDataCountResource(t *testing.T) {
+	m := testModule(t, "plan-computed-data-count")
+	p := testProvider("aws")
+	p.DiffFn = testDiffFn
+	ctx := testContext2(t, &ContextOpts{
+		Module: m,
+		Providers: map[string]ResourceProviderFactory{
+			"aws": testProviderFuncFixed(p),
+		},
+	})
+
+	plan, err := ctx.Plan()
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	if got := len(plan.Diff.Modules); got != 1 {
+		t.Fatalf("got %d modules; want 1", got)
+	}
+
+	moduleDiff := plan.Diff.Modules[0]
+
+	// make sure we created 3 "bar"s
+	for i := 0; i < 3; i++ {
+		resource := fmt.Sprintf("data.aws_vpc.bar.%d", i)
+		if _, ok := moduleDiff.Resources[resource]; !ok {
+			t.Fatalf("missing diff for %s", resource)
+		}
+	}
+}
+
 // Higher level test at TestResource_dataSourceListPlanPanic
 func TestContext2Plan_dataSourceTypeMismatch(t *testing.T) {
 	m := testModule(t, "plan-data-source-type-mismatch")

terraform/test-fixtures/plan-computed-data-count/main.tf

@@ -0,0 +1,9 @@
+resource "aws_instance" "foo" {
+  num     = "2"
+  compute = "foo"
+}
+
+data "aws_vpc" "bar" {
+  count = 3
+  foo = "${aws_instance.foo.foo}"
+}