failing tests when using resources with count

Two different tests failing around resourced with count

builtin/providers/test/data_source.go

@@ -31,6 +31,17 @@ func testDataSource() *schema.Resource {
 				Type:     schema.TypeString,
 				Computed: true,
 			},
+
+			"input_map": {
+				Type:     schema.TypeMap,
+				Elem:     &schema.Schema{Type: schema.TypeString},
+				Optional: true,
+			},
+			"output_map": {
+				Type:     schema.TypeMap,
+				Elem:     &schema.Schema{Type: schema.TypeString},
+				Computed: true,
+			},
 		},
 	}
 }
@@ -45,5 +56,8 @@ func testDataSourceRead(d *schema.ResourceData, meta interface{}) error {
 		d.Set("output", "some output")
 	}
 
+	if inputMap, hasInput := d.GetOk("input_map"); hasInput {
+		d.Set("output_map", inputMap)
+	}
 	return nil
 }

builtin/providers/test/data_source_test.go

@@ -3,6 +3,7 @@ package test
 import (
 	"errors"
 	"fmt"
+	"regexp"
 	"strings"
 	"testing"
 
@@ -189,3 +190,52 @@ data "test_data_source" "y" {
 		},
 	})
 }
+
+// referencing test_data_source.one.output_map["a"] should produce an error when
+// there's a count, but instead we end up with an unknown value after apply.
+func TestDataSource_indexedCountOfOne(t *testing.T) {
+	resource.UnitTest(t, resource.TestCase{
+		Providers: testAccProviders,
+		Steps: []resource.TestStep{
+			{
+				Config: strings.TrimSpace(`
+data "test_data_source" "one" {
+	count = 1
+  input_map = {
+		"a" = "b"
+	}
+}
+
+data "test_data_source" "two" {
+	input_map = {
+		"x" = data.test_data_source.one.output_map["a"]
+	}
+}
+				`),
+				ExpectError: regexp.MustCompile("value does not have any attributes"),
+			},
+		},
+	})
+}
+
+// Verify that we can destroy when a data source references something with a
+// count of 1.
+func TestDataSource_countRefDestroyError(t *testing.T) {
+	resource.UnitTest(t, resource.TestCase{
+		Providers: testAccProviders,
+		Steps: []resource.TestStep{
+			{
+				Config: strings.TrimSpace(`
+data "test_data_source" "one" {
+  count = 1
+  input = "a"
+}
+
+data "test_data_source" "two" {
+  input = data.test_data_source.one[0].output
+}
+				`),
+			},
+		},
+	})
+}

builtin/providers/test/resource_test.go

@@ -493,3 +493,31 @@ resource "test_resource" "foo" {
 		},
 	})
 }
+
+// Verify that we can destroy when a managed resource references something with
+// a count of 1.
+func TestResource_countRefDestroyError(t *testing.T) {
+	resource.UnitTest(t, resource.TestCase{
+		Providers: testAccProviders,
+		Steps: []resource.TestStep{
+			{
+				Config: strings.TrimSpace(`
+resource "test_resource" "one" {
+	count = 1
+	required     = "ok"
+	required_map = {
+	  key = "val"
+	}
+}
+
+resource "test_resource" "two" {
+	required     = test_resource.one[0].id
+	required_map = {
+	  key = "val"
+	}
+}
+				`),
+			},
+		},
+	})
+}