[testing framework] implement expect_failures functionality (#33443)

internal/addrs/check_rule_diagnostic.go

@@ -0,0 +1,70 @@
+package addrs
+
+import "github.com/hashicorp/terraform/internal/tfdiags"
+
+// DiagnosticExtraCheckRule provides an interface for diagnostic ExtraInfo to
+// retrieve an embedded CheckRule from within a tfdiags.Diagnostic.
+type DiagnosticExtraCheckRule interface {
+	// DiagnosticOriginatesFromCheckRule returns the CheckRule that the
+	// surrounding diagnostic originated from.
+	DiagnosticOriginatesFromCheckRule() CheckRule
+}
+
+// DiagnosticOriginatesFromCheckRule checks if the provided diagnostic contains
+// a CheckRule as ExtraInfo and returns that CheckRule and true if it does. This
+// function returns an empty CheckRule and false if the diagnostic does not
+// contain a CheckRule.
+func DiagnosticOriginatesFromCheckRule(diag tfdiags.Diagnostic) (CheckRule, bool) {
+	maybe := tfdiags.ExtraInfo[DiagnosticExtraCheckRule](diag)
+	if maybe == nil {
+		return CheckRule{}, false
+	}
+	return maybe.DiagnosticOriginatesFromCheckRule(), true
+}
+
+// CheckRuleDiagnosticExtra is an object that can be attached to diagnostics
+// that originate from check rules.
+//
+// It implements the DiagnosticExtraCheckRule interface for retrieving the
+// concrete CheckRule that spawned the diagnostic.
+//
+// It also implements the tfdiags.DiagnosticExtraDoNotConsolidate interface, to
+// stop diagnostics created by check blocks being consolidated.
+//
+// It also implements the tfdiags.DiagnosticExtraUnwrapper interface, as nested
+// data blocks will attach this struct but do want to lose any extra info
+// embedded in the original diagnostic.
+type CheckRuleDiagnosticExtra struct {
+	CheckRule CheckRule
+
+	wrapped interface{}
+}
+
+var (
+	_ DiagnosticExtraCheckRule                = (*CheckRuleDiagnosticExtra)(nil)
+	_ tfdiags.DiagnosticExtraDoNotConsolidate = (*CheckRuleDiagnosticExtra)(nil)
+	_ tfdiags.DiagnosticExtraUnwrapper        = (*CheckRuleDiagnosticExtra)(nil)
+	_ tfdiags.DiagnosticExtraWrapper          = (*CheckRuleDiagnosticExtra)(nil)
+)
+
+func (c *CheckRuleDiagnosticExtra) UnwrapDiagnosticExtra() interface{} {
+	return c.wrapped
+}
+
+func (c *CheckRuleDiagnosticExtra) WrapDiagnosticExtra(inner interface{}) {
+	if c.wrapped != nil {
+		// This is a logical inconsistency, the caller should know whether they
+		// have already wrapped an extra or not.
+		panic("Attempted to wrap a diagnostic extra into a CheckRuleDiagnosticExtra that is already wrapping a different extra. This is a bug in Terraform, please report it.")
+	}
+	c.wrapped = inner
+}
+
+func (c *CheckRuleDiagnosticExtra) DoNotConsolidateDiagnostic() bool {
+	// Do not consolidate warnings from check blocks.
+	return c.CheckRule.Container.CheckableKind() == CheckableCheck
+}
+
+func (c *CheckRuleDiagnosticExtra) DiagnosticOriginatesFromCheckRule() CheckRule {
+	return c.CheckRule
+}

internal/addrs/check_rule_diagnostic_test.go

@@ -0,0 +1,108 @@
+package addrs
+
+import (
+	"testing"
+
+	"github.com/hashicorp/hcl/v2"
+
+	"github.com/hashicorp/terraform/internal/tfdiags"
+)
+
+func TestCheckRuleDiagnosticExtra_WrapsExtra(t *testing.T) {
+	var originals tfdiags.Diagnostics
+	originals = originals.Append(&hcl.Diagnostic{
+		Severity: hcl.DiagError,
+		Summary:  "original error",
+		Detail:   "this is an error",
+		Extra:    "extra",
+	})
+
+	overridden := tfdiags.OverrideAll(originals, tfdiags.Warning, func() tfdiags.DiagnosticExtraWrapper {
+		return &CheckRuleDiagnosticExtra{}
+	})
+
+	if overridden[0].ExtraInfo().(*CheckRuleDiagnosticExtra).wrapped.(string) != "extra" {
+		t.Errorf("unexpected extra info: %v", overridden[0].ExtraInfo())
+	}
+}
+
+func TestCheckRuleDiagnosticExtra_Unwraps(t *testing.T) {
+	var originals tfdiags.Diagnostics
+	originals = originals.Append(&hcl.Diagnostic{
+		Severity: hcl.DiagError,
+		Summary:  "original error",
+		Detail:   "this is an error",
+		Extra:    "extra",
+	})
+
+	overridden := tfdiags.OverrideAll(originals, tfdiags.Warning, func() tfdiags.DiagnosticExtraWrapper {
+		return &CheckRuleDiagnosticExtra{}
+	})
+
+	result := tfdiags.ExtraInfo[string](overridden[0])
+	if result != "extra" {
+		t.Errorf("unexpected extra info: %v", result)
+	}
+}
+
+func TestCheckRuleDiagnosticExtra_DoNotConsolidate(t *testing.T) {
+	var diags tfdiags.Diagnostics
+	diags = diags.Append(&hcl.Diagnostic{
+		Severity: hcl.DiagError,
+		Summary:  "original error",
+		Detail:   "this is an error",
+		Extra: &CheckRuleDiagnosticExtra{
+			CheckRule: NewCheckRule(AbsOutputValue{
+				Module: RootModuleInstance,
+				OutputValue: OutputValue{
+					Name: "output",
+				},
+			}, OutputPrecondition, 0),
+		},
+	})
+	diags = diags.Append(&hcl.Diagnostic{
+		Severity: hcl.DiagError,
+		Summary:  "original error",
+		Detail:   "this is an error",
+		Extra: &CheckRuleDiagnosticExtra{
+			CheckRule: NewCheckRule(AbsCheck{
+				Module: RootModuleInstance,
+				Check: Check{
+					Name: "check",
+				},
+			}, CheckAssertion, 0),
+		},
+	})
+
+	if tfdiags.DoNotConsolidateDiagnostic(diags[0]) {
+		t.Errorf("first diag should be consolidated but was not")
+	}
+
+	if !tfdiags.DoNotConsolidateDiagnostic(diags[1]) {
+		t.Errorf("second diag should not be consolidated but was")
+	}
+
+}
+
+func TestDiagnosticOriginatesFromCheckRule_Passes(t *testing.T) {
+	var diags tfdiags.Diagnostics
+	diags = diags.Append(&hcl.Diagnostic{
+		Severity: hcl.DiagError,
+		Summary:  "original error",
+		Detail:   "this is an error",
+	})
+	diags = diags.Append(&hcl.Diagnostic{
+		Severity: hcl.DiagError,
+		Summary:  "original error",
+		Detail:   "this is an error",
+		Extra:    &CheckRuleDiagnosticExtra{},
+	})
+
+	if _, ok := DiagnosticOriginatesFromCheckRule(diags[0]); ok {
+		t.Errorf("first diag did not originate from check rule but thinks it did")
+	}
+
+	if _, ok := DiagnosticOriginatesFromCheckRule(diags[1]); !ok {
+		t.Errorf("second diag did originate from check rule but this it did not")
+	}
+}

internal/backend/local/backend_apply.go

@@ -10,6 +10,7 @@ import (
 	"log"
 	"time"
 
+	"github.com/hashicorp/terraform/internal/addrs"
 	"github.com/hashicorp/terraform/internal/backend"
 	"github.com/hashicorp/terraform/internal/command/views"
 	"github.com/hashicorp/terraform/internal/logging"
@@ -198,9 +199,10 @@ func (b *Local) opApply(
 			// is needlessly confusing.
 			var filteredDiags tfdiags.Diagnostics
 			for _, diag := range diags {
-				if !tfdiags.IsFromCheckBlock(diag) {
-					filteredDiags = filteredDiags.Append(diag)
+				if rule, ok := addrs.DiagnosticOriginatesFromCheckRule(diag); ok && rule.Container.CheckableKind() == addrs.CheckableCheck {
+					continue
 				}
+				filteredDiags = filteredDiags.Append(diag)
 			}
 			diags = filteredDiags
 		}

internal/configs/parser_config_dir_test.go

@@ -116,6 +116,7 @@ func TestParserLoadConfigDirSuccess(t *testing.T) {
 func TestParserLoadConfigDirWithTests(t *testing.T) {
 	directories := []string{
 		"testdata/valid-modules/with-tests",
+		"testdata/valid-modules/with-tests-expect-failures",
 		"testdata/valid-modules/with-tests-nested",
 		"testdata/valid-modules/with-tests-json",
 	}

internal/configs/test_file.go

@@ -78,6 +78,11 @@ type TestRun struct {
 	// checked by this run block.
 	CheckRules []*CheckRule
 
+	// ExpectFailures should be a list of checkable objects that are expected
+	// to report a failure from their custom conditions as part of this test
+	// run.
+	ExpectFailures []hcl.Traversal
+
 	NameDeclRange      hcl.Range
 	VariablesDeclRange hcl.Range
 	DeclRange          hcl.Range
@@ -233,6 +238,12 @@ func decodeTestRunBlock(block *hcl.Block) (*TestRun, hcl.Diagnostics) {
 		r.Command = ApplyTestCommand // Default to apply
 	}
 
+	if attr, exists := content.Attributes["expect_failures"]; exists {
+		failures, failDiags := decodeDependsOn(attr)
+		diags = append(diags, failDiags...)
+		r.ExpectFailures = failures
+	}
+
 	return &r, diags
 }
 
@@ -311,6 +322,7 @@ var testFileSchema = &hcl.BodySchema{
 var testRunBlockSchema = &hcl.BodySchema{
 	Attributes: []hcl.AttributeSchema{
 		{Name: "command"},
+		{Name: "expect_failures"},
 	},
 	Blocks: []hcl.BlockHeaderSchema{
 		{

internal/configs/testdata/valid-modules/with-tests-expect-failures/main.tf

@@ -0,0 +1,13 @@
+
+variable "input" {
+  type = string
+}
+
+
+resource "foo_resource" "a" {
+  value = var.input
+}
+
+output "output" {
+  value = foo_resource.a.value
+}

internal/configs/testdata/valid-modules/with-tests-expect-failures/test_case_one.tftest

@@ -0,0 +1,10 @@
+variables {
+  input = "default"
+}
+
+run "test_run_one" {
+  expect_failures = [
+    input.input,
+    output.output,
+  ]
+}

internal/configs/testdata/valid-modules/with-tests-expect-failures/test_case_two.tftest

@@ -0,0 +1,9 @@
+variables {
+  input = "default"
+}
+
+run "test_run_one" {
+  expect_failures = [
+    foo_resource.a,
+  ]
+}

internal/moduletest/run.go

@@ -1,6 +1,11 @@
 package moduletest
 
 import (
+	"fmt"
+
+	"github.com/hashicorp/hcl/v2"
+
+	"github.com/hashicorp/terraform/internal/addrs"
 	"github.com/hashicorp/terraform/internal/configs"
 	"github.com/hashicorp/terraform/internal/tfdiags"
 )
@@ -13,3 +18,208 @@ type Run struct {
 
 	Diagnostics tfdiags.Diagnostics
 }
+
+// ValidateExpectedFailures steps through the provided diagnostics (which should
+// be the result of a plan or an apply operation), and does 3 things:
+//  1. Removes diagnostics that match the expected failures from the config.
+//  2. Upgrades warnings from check blocks into errors where appropriate so the
+//     test will fail later.
+//  3. Adds diagnostics for any expected failures that were not satisfied.
+//
+// Point 2 is a bit complicated so worth expanding on. In normal Terraform
+// execution, any error that originates within a check block (either from an
+// assertion or a scoped data source) is wrapped up as a Warning to be
+// identified to the user but not to fail the actual Terraform operation. During
+// test execution, we want to upgrade (or rollback) these warnings into errors
+// again so the test will fail. We do that as part of this function as we are
+// already processing the diagnostics from check blocks in here anyway.
+//
+// The way the function works out which diagnostics are relevant to expected
+// failures is by using the tfdiags.ValuedDiagnostic functionality to detect
+// which diagnostics were generated by custom conditions. Terraform adds the
+// addrs.CheckRule that generated each diagnostic to the diagnostic itself so we
+// can tell which diagnostics can be expected.
+func (run *Run) ValidateExpectedFailures(originals tfdiags.Diagnostics) tfdiags.Diagnostics {
+
+	// We're going to capture all the checkable objects that are referenced
+	// from the expected failures.
+	expectedFailures := addrs.MakeMap[addrs.Referenceable, bool]()
+	sourceRanges := addrs.MakeMap[addrs.Referenceable, tfdiags.SourceRange]()
+
+	for _, traversal := range run.Config.ExpectFailures {
+		// Ignore the diagnostics returned from the reference parsing, these
+		// references will have been checked earlier in the process by the
+		// validate stage so we don't need to do that again here.
+		reference, _ := addrs.ParseRefFromTestingScope(traversal)
+		expectedFailures.Put(reference.Subject, false)
+		sourceRanges.Put(reference.Subject, reference.SourceRange)
+	}
+
+	var diags tfdiags.Diagnostics
+	for _, diag := range originals {
+
+		if rule, ok := addrs.DiagnosticOriginatesFromCheckRule(diag); ok {
+			switch rule.Container.CheckableKind() {
+			case addrs.CheckableOutputValue:
+				addr := rule.Container.(addrs.AbsOutputValue)
+				if !addr.Module.IsRoot() {
+					// failures can only be expected against checkable objects
+					// in the root module. This diagnostic will be added into
+					// returned set below.
+					break
+				}
+
+				if diag.Severity() == tfdiags.Warning {
+					// Warnings don't count as errors. This diagnostic will be
+					// added into the returned set below.
+					break
+				}
+
+				if expectedFailures.Has(addr.OutputValue) {
+					// Then this failure is expected! Mark the original map as
+					// having found a failure and swallow this error by
+					// continuing and not adding it into the returned set of
+					// diagnostics.
+					expectedFailures.Put(addr.OutputValue, true)
+					continue
+				}
+
+				// Otherwise, this isn't an expected failure so just fall out
+				// and add it into the returned set of diagnostics below.
+
+			case addrs.CheckableResource:
+				addr := rule.Container.(addrs.AbsResourceInstance)
+				if !addr.Module.IsRoot() {
+					// failures can only be expected against checkable objects
+					// in the root module. This diagnostic will be added into
+					// returned set below.
+					break
+				}
+
+				if diag.Severity() == tfdiags.Warning {
+					// Warnings don't count as errors. This diagnostic will be
+					// added into the returned set below.
+					break
+				}
+
+				if expectedFailures.Has(addr.Resource) {
+					// Then this failure is expected! Mark the original map as
+					// having found a failure and swallow this error by
+					// continuing and not adding it into the returned set of
+					// diagnostics.
+					expectedFailures.Put(addr.Resource, true)
+					continue
+				}
+
+				if expectedFailures.Has(addr.Resource.Resource) {
+					// We can also blanket expect failures in all instances for
+					// a resource so we check for that here as well.
+					expectedFailures.Put(addr.Resource.Resource, true)
+					continue
+				}
+
+				// Otherwise, this isn't an expected failure so just fall out
+				// and add it into the returned set of diagnostics below.
+
+			case addrs.CheckableCheck:
+				addr := rule.Container.(addrs.AbsCheck)
+
+				// Check blocks are a bit more difficult than the others. Check
+				// block diagnostics could be from a nested data block, or
+				// from a failed assertion, and have all been marked as just
+				// warning severity.
+				//
+				// For diagnostics from failed assertions, we want to check if
+				// it was expected and skip it if it was. But if it wasn't
+				// expected we want to upgrade the diagnostic from a warning
+				// into an error so the test case will fail overall.
+				//
+				// For diagnostics from nested data blocks, we have two
+				// categories of diagnostics. First, diagnostics that were
+				// originally errors and we mapped into warnings. Second,
+				// diagnostics that were originally warnings and stayed that
+				// way. For the first case, we want to turn these back to errors
+				// and use them as part of the expected failures functionality.
+				// The second case should remain as warnings and be ignored by
+				// the expected failures functionality.
+				//
+				// Note, as well that we still want to upgrade failed checks
+				// from child modules into errors, so in the other branches we
+				// just do a simple blanket skip off all diagnostics not
+				// from the root module. We're more selective here, only
+				// diagnostics from the root module are considered for the
+				// expect failures functionality but we do also upgrade
+				// diagnostics from child modules back into errors.
+
+				if rule.Type == addrs.CheckAssertion {
+					// Then this diagnostic is from a check block assertion, it
+					// is something we want to treat as an error even though it
+					// is actually claiming to be a warning.
+
+					if addr.Module.IsRoot() && expectedFailures.Has(addr.Check) {
+						// Then this failure is expected! Mark the original map as
+						// having found a failure and continue.
+						expectedFailures.Put(addr.Check, true)
+						continue
+					}
+
+					// Otherwise, let's package this up as an error and move on.
+					diags = diags.Append(tfdiags.Override(diag, tfdiags.Error, nil))
+					continue
+				} else if rule.Type == addrs.CheckDataResource {
+					// Then the diagnostic we have was actually overridden so
+					// let's get back to the original.
+					original := tfdiags.UndoOverride(diag)
+
+					// This diagnostic originated from a scoped data source.
+					if addr.Module.IsRoot() && original.Severity() == tfdiags.Error {
+						// Okay, we have a genuine error from the root module,
+						// so we can now check if we want to ignore it or not.
+						if expectedFailures.Has(addr.Check) {
+							// Then this failure is expected! Mark the original map as
+							// having found a failure and continue.
+							expectedFailures.Put(addr.Check, true)
+							continue
+						}
+					}
+
+					// In all other cases, we want to add the original error
+					// into the set we return to the testing framework and move
+					// onto the next one.
+					diags = diags.Append(original)
+					continue
+				} else {
+					panic("invalid CheckType: " + rule.Type.String())
+				}
+			default:
+				panic("unrecognized CheckableKind: " + rule.Container.CheckableKind().String())
+			}
+		}
+
+		// If we get here, then we're not modifying the original diagnostic at
+		// all. We just want the testing framework to treat it as normal.
+		diags = diags.Append(diag)
+	}
+
+	// Okay, we've checked all our diagnostics to see if any were expected.
+	// Now, let's make sure that all the checkable objects we expected to fail
+	// actually did!
+
+	for _, elem := range expectedFailures.Elems {
+		addr := elem.Key
+		failed := elem.Value
+
+		if !failed {
+			// Then we expected a failure, and it did not occur. Add it to the
+			// diagnostics.
+			diags = diags.Append(&hcl.Diagnostic{
+				Severity: hcl.DiagError,
+				Summary:  "Missing expected failure",
+				Detail:   fmt.Sprintf("The checkable object, %s, was expected to report an error but did not.", addr.String()),
+				Subject:  sourceRanges.Get(addr).ToHCL().Ptr(),
+			})
+		}
+	}
+
+	return diags
+}

internal/moduletest/run_test.go

@@ -0,0 +1,692 @@
+package moduletest
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/hcl/v2/hclsyntax"
+
+	"github.com/hashicorp/terraform/internal/addrs"
+	"github.com/hashicorp/terraform/internal/configs"
+	"github.com/hashicorp/terraform/internal/tfdiags"
+)
+
+func TestRun_ValidateExpectedFailures(t *testing.T) {
+
+	type output struct {
+		Description tfdiags.Description
+		Severity    tfdiags.Severity
+	}
+
+	tcs := map[string]struct {
+		ExpectedFailures []string
+		Input            tfdiags.Diagnostics
+		Output           []output
+	}{
+		"empty": {
+			ExpectedFailures: nil,
+			Input:            nil,
+			Output:           nil,
+		},
+		"carries through simple diags": {
+			Input: createDiagnostics(func(diags tfdiags.Diagnostics) tfdiags.Diagnostics {
+
+				diags = diags.Append(&hcl.Diagnostic{
+					Severity: hcl.DiagError,
+					Summary:  "simple error",
+					Detail:   "want to see this in the returned set",
+				})
+
+				diags = diags.Append(&hcl.Diagnostic{
+					Severity: hcl.DiagWarning,
+					Summary:  "simple warning",
+					Detail:   "want to see this in the returned set",
+				})
+
+				return diags
+			}),
+			Output: []output{
+				{
+					Description: tfdiags.Description{
+						Summary: "simple error",
+						Detail:  "want to see this in the returned set",
+					},
+					Severity: tfdiags.Error,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "simple warning",
+						Detail:  "want to see this in the returned set",
+					},
+					Severity: tfdiags.Warning,
+				},
+			},
+		},
+		"expected failures did not fail": {
+			ExpectedFailures: []string{
+				"check.example",
+			},
+			Input: nil,
+			Output: []output{
+				{
+					Description: tfdiags.Description{
+						Summary: "Missing expected failure",
+						Detail:  "The checkable object, check.example, was expected to report an error but did not.",
+					},
+					Severity: tfdiags.Error,
+				},
+			},
+		},
+		"outputs": {
+			ExpectedFailures: []string{
+				"output.expected_one",
+				"output.expected_two",
+			},
+			Input: createDiagnostics(func(diags tfdiags.Diagnostics) tfdiags.Diagnostics {
+
+				// First, let's create an output that failed that isn't
+				// expected. This should be unaffected by our function.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "unexpected failure",
+						Detail:   "this should not be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsOutputValue{
+								Module:      addrs.RootModuleInstance,
+								OutputValue: addrs.OutputValue{Name: "unexpected"},
+							}, addrs.OutputPrecondition, 0),
+						},
+					})
+
+				// Second, let's create an output that failed but is expected.
+				// Our function should remove this from the set of diags.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "expected failure",
+						Detail:   "this should be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsOutputValue{
+								Module:      addrs.RootModuleInstance,
+								OutputValue: addrs.OutputValue{Name: "expected_one"},
+							}, addrs.OutputPrecondition, 0),
+						},
+					})
+
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagWarning,
+						Summary:  "expected warning",
+						Detail:   "this should not be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsOutputValue{
+								Module:      addrs.RootModuleInstance,
+								OutputValue: addrs.OutputValue{Name: "expected_one"},
+							}, addrs.OutputPrecondition, 0),
+						},
+					})
+
+				// The error we are adding here is for expected_two but in a
+				// child module. We expect that this diagnostic shouldn't
+				// trigger our expected failure, and that an extra diagnostic
+				// should be created complaining that the output wasn't actually
+				// triggered.
+
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "error in child module",
+						Detail:   "this should not be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsOutputValue{
+								Module: []addrs.ModuleInstanceStep{
+									{
+										Name: "child_module",
+									},
+								},
+								OutputValue: addrs.OutputValue{Name: "expected_two"},
+							}, addrs.OutputPrecondition, 0),
+						},
+					})
+
+				return diags
+			}),
+			Output: []output{
+				{
+					Description: tfdiags.Description{
+						Summary: "unexpected failure",
+						Detail:  "this should not be removed",
+					},
+					Severity: tfdiags.Error,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "expected warning",
+						Detail:  "this should not be removed",
+					},
+					Severity: tfdiags.Warning,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "error in child module",
+						Detail:  "this should not be removed",
+					},
+					Severity: tfdiags.Error,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "Missing expected failure",
+						Detail:  "The checkable object, output.expected_two, was expected to report an error but did not.",
+					},
+					Severity: tfdiags.Error,
+				},
+			},
+		},
+		"resources": {
+			ExpectedFailures: []string{
+				"test_instance.single",
+				"test_instance.all_instances",
+				"test_instance.instance[0]",
+				"test_instance.instance[2]",
+				"test_instance.missing",
+			},
+			Input: createDiagnostics(func(diags tfdiags.Diagnostics) tfdiags.Diagnostics {
+				// First, we'll create an unexpected failure that should be
+				// carried through untouched.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "unexpected failure",
+						Detail:   "this should not be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: addrs.RootModuleInstance,
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "unexpected",
+									},
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+
+				// Second, we'll create a failure from our test_instance.single
+				// resource that should be removed.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "expected failure in test_instance.single",
+						Detail:   "this should be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: addrs.RootModuleInstance,
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "single",
+									},
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+
+				// Third, we'll create a warning from our test_instance.single
+				// resource that should be propagated as it is only a warning.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagWarning,
+						Summary:  "expected warning in test_instance.single",
+						Detail:   "this should not be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: addrs.RootModuleInstance,
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "single",
+									},
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+
+				// Fourth, we'll create diagnostics from several instances of
+				// the test_instance.all_instances which should all be removed.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "expected failure in test_instance.all_instances[0]",
+						Detail:   "this should be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: addrs.RootModuleInstance,
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "all_instances",
+									},
+									Key: addrs.IntKey(0),
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "expected failure in test_instance.all_instances[1]",
+						Detail:   "this should be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: addrs.RootModuleInstance,
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "all_instances",
+									},
+									Key: addrs.IntKey(1),
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "expected failure in test_instance.all_instances[2]",
+						Detail:   "this should be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: addrs.RootModuleInstance,
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "all_instances",
+									},
+									Key: addrs.IntKey(2),
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+
+				// Fifth, we'll create diagnostics for several instances of
+				// the test_instance.instance resource, only some of which
+				// should be removed.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "expected failure in test_instance.instance[0]",
+						Detail:   "this should be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: addrs.RootModuleInstance,
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "instance",
+									},
+									Key: addrs.IntKey(0),
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "expected failure in test_instance.instance[1]",
+						Detail:   "this should not be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: addrs.RootModuleInstance,
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "instance",
+									},
+									Key: addrs.IntKey(1),
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "expected failure in test_instance.instance[2]",
+						Detail:   "this should be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: addrs.RootModuleInstance,
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "instance",
+									},
+									Key: addrs.IntKey(2),
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+
+				// Finally, we'll create an error that originated from
+				// test_instance.missing but in a child module which shouldn't
+				// be removed.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagError,
+						Summary:  "failure in child module",
+						Detail:   "this should not be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsResourceInstance{
+								Module: []addrs.ModuleInstanceStep{
+									{
+										Name: "child_module",
+									},
+								},
+								Resource: addrs.ResourceInstance{
+									Resource: addrs.Resource{
+										Mode: addrs.ManagedResourceMode,
+										Type: "test_instance",
+										Name: "missing",
+									},
+								},
+							}, addrs.ResourcePrecondition, 0),
+						},
+					})
+
+				return diags
+			}),
+			Output: []output{
+				{
+					Description: tfdiags.Description{
+						Summary: "unexpected failure",
+						Detail:  "this should not be removed",
+					},
+					Severity: tfdiags.Error,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "expected warning in test_instance.single",
+						Detail:  "this should not be removed",
+					},
+					Severity: tfdiags.Warning,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "expected failure in test_instance.instance[1]",
+						Detail:  "this should not be removed",
+					},
+					Severity: tfdiags.Error,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "failure in child module",
+						Detail:  "this should not be removed",
+					},
+					Severity: tfdiags.Error,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "Missing expected failure",
+						Detail:  "The checkable object, test_instance.missing, was expected to report an error but did not.",
+					},
+					Severity: tfdiags.Error,
+				},
+			},
+		},
+		"check_assertions": {
+			ExpectedFailures: []string{
+				"check.expected",
+				"check.missing",
+			},
+			Input: createDiagnostics(func(diags tfdiags.Diagnostics) tfdiags.Diagnostics {
+				// First, we'll add an unexpected warning from a check block
+				// assertion that should get upgraded to an error.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagWarning,
+						Summary:  "unexpected failure",
+						Detail:   "this should upgrade and not be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsCheck{
+								Module: addrs.RootModuleInstance,
+								Check: addrs.Check{
+									Name: "unexpected",
+								},
+							}, addrs.CheckAssertion, 0),
+						},
+					})
+
+				// Second, we'll add an unexpected warning from a check block
+				// in a child module that should get upgrade to error.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagWarning,
+						Summary:  "expected failure in child module",
+						Detail:   "this should upgrade and not be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsCheck{
+								Module: []addrs.ModuleInstanceStep{
+									{
+										Name: "child_module",
+									},
+								},
+								Check: addrs.Check{
+									Name: "expected",
+								},
+							}, addrs.CheckAssertion, 0),
+						},
+					})
+
+				// Third, we'll add an expected warning from a check block
+				// assertion that should be removed.
+				diags = diags.Append(
+					&hcl.Diagnostic{
+						Severity: hcl.DiagWarning,
+						Summary:  "expected failure",
+						Detail:   "this should be removed",
+						Extra: &addrs.CheckRuleDiagnosticExtra{
+							CheckRule: addrs.NewCheckRule(addrs.AbsCheck{
+								Module: addrs.RootModuleInstance,
+								Check: addrs.Check{
+									Name: "expected",
+								},
+							}, addrs.CheckAssertion, 0),
+						},
+					})
+
+				// The second expected failure has no diagnostics, we just want
+				// to make sure that a new diagnostic is added for this case.
+
+				return diags
+			}),
+			Output: []output{
+				{
+					Description: tfdiags.Description{
+						Summary: "unexpected failure",
+						Detail:  "this should upgrade and not be removed",
+					},
+					Severity: tfdiags.Error,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "expected failure in child module",
+						Detail:  "this should upgrade and not be removed",
+					},
+					Severity: tfdiags.Error,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "Missing expected failure",
+						Detail:  "The checkable object, check.missing, was expected to report an error but did not.",
+					},
+					Severity: tfdiags.Error,
+				},
+			},
+		},
+		"check_data_sources": {
+			ExpectedFailures: []string{
+				"check.expected",
+			},
+			Input: createDiagnostics(func(diags tfdiags.Diagnostics) tfdiags.Diagnostics {
+				// First, we'll add an unexpected warning from a check block
+				// assertion that should be propagated as an error.
+				diags = diags.Append(
+					tfdiags.Override(
+						tfdiags.Sourceless(tfdiags.Error, "unexpected failure", "this should be an error and not removed"),
+						tfdiags.Warning,
+						func() tfdiags.DiagnosticExtraWrapper {
+							return &addrs.CheckRuleDiagnosticExtra{
+								CheckRule: addrs.NewCheckRule(addrs.AbsCheck{
+									Module: addrs.RootModuleInstance,
+									Check: addrs.Check{
+										Name: "unexpected",
+									},
+								}, addrs.CheckDataResource, 0),
+							}
+						}))
+
+				// Second, we'll add an unexpected warning from a check block
+				// assertion that should remain as a warning.
+				diags = diags.Append(
+					tfdiags.Override(
+						tfdiags.Sourceless(tfdiags.Warning, "unexpected warning", "this should be a warning and not removed"),
+						tfdiags.Warning,
+						func() tfdiags.DiagnosticExtraWrapper {
+							return &addrs.CheckRuleDiagnosticExtra{
+								CheckRule: addrs.NewCheckRule(addrs.AbsCheck{
+									Module: addrs.RootModuleInstance,
+									Check: addrs.Check{
+										Name: "unexpected",
+									},
+								}, addrs.CheckDataResource, 0),
+							}
+						}))
+
+				// Third, we'll add an unexpected warning from a check block
+				// in a child module that should be propagated as an error.
+				diags = diags.Append(
+					tfdiags.Override(
+						tfdiags.Sourceless(tfdiags.Error, "expected failure from child module", "this should be an error and not removed"),
+						tfdiags.Warning,
+						func() tfdiags.DiagnosticExtraWrapper {
+							return &addrs.CheckRuleDiagnosticExtra{
+								CheckRule: addrs.NewCheckRule(addrs.AbsCheck{
+									Module: []addrs.ModuleInstanceStep{
+										{
+											Name: "child_module",
+										},
+									},
+									Check: addrs.Check{
+										Name: "expected",
+									},
+								}, addrs.CheckDataResource, 0),
+							}
+						}))
+
+				// Fourth, we'll add an expected warning that should be removed.
+				diags = diags.Append(
+					tfdiags.Override(
+						tfdiags.Sourceless(tfdiags.Error, "expected failure", "this should be removed"),
+						tfdiags.Warning,
+						func() tfdiags.DiagnosticExtraWrapper {
+							return &addrs.CheckRuleDiagnosticExtra{
+								CheckRule: addrs.NewCheckRule(addrs.AbsCheck{
+									Module: addrs.RootModuleInstance,
+									Check: addrs.Check{
+										Name: "expected",
+									},
+								}, addrs.CheckDataResource, 0),
+							}
+						}))
+
+				return diags
+			}),
+			Output: []output{
+				{
+					Description: tfdiags.Description{
+						Summary: "unexpected failure",
+						Detail:  "this should be an error and not removed",
+					},
+					Severity: tfdiags.Error,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "unexpected warning",
+						Detail:  "this should be a warning and not removed",
+					},
+					Severity: tfdiags.Warning,
+				},
+				{
+					Description: tfdiags.Description{
+						Summary: "expected failure from child module",
+						Detail:  "this should be an error and not removed",
+					},
+					Severity: tfdiags.Error,
+				},
+			},
+		},
+	}
+	for name, tc := range tcs {
+		t.Run(name, func(t *testing.T) {
+			var traversals []hcl.Traversal
+			for _, ef := range tc.ExpectedFailures {
+				traversal, diags := hclsyntax.ParseTraversalAbs([]byte(ef), "foo.tf", hcl.Pos{Line: 1, Column: 1})
+				if diags.HasErrors() {
+					t.Errorf("invalid expected failure %s: %v", ef, diags.Error())
+				}
+				traversals = append(traversals, traversal)
+			}
+
+			if t.Failed() {
+				return
+			}
+
+			run := Run{
+				Config: &configs.TestRun{
+					ExpectFailures: traversals,
+				},
+			}
+
+			out := run.ValidateExpectedFailures(tc.Input)
+			ix := 0
+			for ; ix < len(tc.Output); ix++ {
+				expected := tc.Output[ix]
+
+				if ix >= len(out) {
+					t.Errorf("missing diagnostic at %d, expected: [%s] %s, %s", ix, expected.Severity, expected.Description.Summary, expected.Description.Detail)
+					continue
+				}
+
+				actual := output{
+					Description: out[ix].Description(),
+					Severity:    out[ix].Severity(),
+				}
+
+				if diff := cmp.Diff(expected, actual); len(diff) > 0 {
+					t.Errorf("mismatched diagnostic at %d:\n%s", ix, diff)
+				}
+			}
+
+			for ; ix < len(out); ix++ {
+				actual := out[ix]
+				t.Errorf("additional diagnostic at %d: [%s] %s, %s", ix, actual.Severity(), actual.Description().Summary, actual.Description().Detail)
+			}
+		})
+	}
+}
+
+func createDiagnostics(populate func(diags tfdiags.Diagnostics) tfdiags.Diagnostics) tfdiags.Diagnostics {
+	var diags tfdiags.Diagnostics
+	diags = populate(diags)
+	return diags
+}

internal/terraform/eval_conditions.go

@@ -51,7 +51,7 @@ func evalCheckRules(typ addrs.CheckRuleType, rules []*configs.CheckRule, ctx Eva
 	severity := diagSeverity.ToHCL()
 
 	for i, rule := range rules {
-		result, ruleDiags := evalCheckRule(typ, rule, ctx, self, keyData, severity)
+		result, ruleDiags := evalCheckRule(addrs.NewCheckRule(self, typ, i), rule, ctx, keyData, severity)
 		diags = diags.Append(ruleDiags)
 
 		log.Printf("[TRACE] evalCheckRules: %s status is now %s", self, result.Status)
@@ -70,7 +70,7 @@ type checkResult struct {
 	FailureMessage string
 }
 
-func validateCheckRule(typ addrs.CheckRuleType, rule *configs.CheckRule, ctx EvalContext, self addrs.Checkable, keyData instances.RepetitionData) (string, *hcl.EvalContext, tfdiags.Diagnostics) {
+func validateCheckRule(addr addrs.CheckRule, rule *configs.CheckRule, ctx EvalContext, keyData instances.RepetitionData) (string, *hcl.EvalContext, tfdiags.Diagnostics) {
 	var diags tfdiags.Diagnostics
 
 	refs, moreDiags := lang.ReferencesInExpr(addrs.ParseRef, rule.Condition)
@@ -80,23 +80,23 @@ func validateCheckRule(typ addrs.CheckRuleType, rule *configs.CheckRule, ctx Eva
 	refs = append(refs, moreRefs...)
 
 	var selfReference, sourceReference addrs.Referenceable
-	switch typ {
+	switch addr.Type {
 	case addrs.ResourcePostcondition:
-		switch s := self.(type) {
+		switch s := addr.Container.(type) {
 		case addrs.AbsResourceInstance:
 			// Only resource postconditions can refer to self
 			selfReference = s.Resource
 		default:
-			panic(fmt.Sprintf("Invalid self reference type %t", self))
+			panic(fmt.Sprintf("Invalid self reference type %t", addr.Container))
 		}
 	case addrs.CheckAssertion:
-		switch s := self.(type) {
+		switch s := addr.Container.(type) {
 		case addrs.AbsCheck:
 			// Only check blocks have scoped resources so need to specify their
 			// source.
 			sourceReference = s.Check
 		default:
-			panic(fmt.Sprintf("Invalid source reference type %t", self))
+			panic(fmt.Sprintf("Invalid source reference type %t", addr.Container))
 		}
 	}
 	scope := ctx.EvaluationScope(selfReference, sourceReference, keyData)
@@ -110,11 +110,11 @@ func validateCheckRule(typ addrs.CheckRuleType, rule *configs.CheckRule, ctx Eva
 	return errorMessage, hclCtx, diags
 }
 
-func evalCheckRule(typ addrs.CheckRuleType, rule *configs.CheckRule, ctx EvalContext, self addrs.Checkable, keyData instances.RepetitionData, severity hcl.DiagnosticSeverity) (checkResult, tfdiags.Diagnostics) {
+func evalCheckRule(addr addrs.CheckRule, rule *configs.CheckRule, ctx EvalContext, keyData instances.RepetitionData, severity hcl.DiagnosticSeverity) (checkResult, tfdiags.Diagnostics) {
 	// NOTE: Intentionally not passing the caller's selected severity in here,
 	// because this reports errors in the configuration itself, not the failure
 	// of an otherwise-valid condition.
-	errorMessage, hclCtx, diags := validateCheckRule(typ, rule, ctx, self, keyData)
+	errorMessage, hclCtx, diags := validateCheckRule(addr, rule, ctx, keyData)
 
 	const errInvalidCondition = "Invalid condition result"
 
@@ -122,22 +122,25 @@ func evalCheckRule(typ addrs.CheckRuleType, rule *configs.CheckRule, ctx EvalCon
 	diags = diags.Append(hclDiags)
 
 	if diags.HasErrors() {
-		log.Printf("[TRACE] evalCheckRule: %s: %s", typ, diags.Err().Error())
+		log.Printf("[TRACE] evalCheckRule: %s: %s", addr.Type, diags.Err().Error())
 		return checkResult{Status: checks.StatusError}, diags
 	}
 
 	if !resultVal.IsKnown() {
 
 		// Check assertions warn if a status is unknown.
-		if typ == addrs.CheckAssertion {
-			diags = diags.Append(tfdiags.AsCheckBlockDiagnostic(&hcl.Diagnostic{
+		if addr.Type == addrs.CheckAssertion {
+			diags = diags.Append(&hcl.Diagnostic{
 				Severity:    hcl.DiagWarning,
-				Summary:     fmt.Sprintf("%s known after apply", typ.Description()),
+				Summary:     fmt.Sprintf("%s known after apply", addr.Type.Description()),
 				Detail:      "The condition could not be evaluated at this time, a result will be known when this plan is applied.",
 				Subject:     rule.Condition.Range().Ptr(),
 				Expression:  rule.Condition,
 				EvalContext: hclCtx,
-			}))
+				Extra: &addrs.CheckRuleDiagnosticExtra{
+					CheckRule: addr,
+				},
+			})
 		}
 
 		// We'll wait until we've learned more, then.
@@ -189,23 +192,20 @@ func evalCheckRule(typ addrs.CheckRuleType, rule *configs.CheckRule, ctx EvalCon
 	if errorMessageForDiags == "" {
 		errorMessageForDiags = "This check failed, but has an invalid error message as described in the other accompanying messages."
 	}
-	diag := &hcl.Diagnostic{
+	diags = diags.Append(&hcl.Diagnostic{
 		// The caller gets to choose the severity of this one, because we
 		// treat condition failures as warnings in the presence of
 		// certain special planning options.
 		Severity:    severity,
-		Summary:     fmt.Sprintf("%s failed", typ.Description()),
+		Summary:     fmt.Sprintf("%s failed", addr.Type.Description()),
 		Detail:      errorMessageForDiags,
 		Subject:     rule.Condition.Range().Ptr(),
 		Expression:  rule.Condition,
 		EvalContext: hclCtx,
-	}
-
-	if typ == addrs.CheckAssertion {
-		diags = diags.Append(tfdiags.AsCheckBlockDiagnostic(diag))
-	} else {
-		diags = diags.Append(diag)
-	}
+		Extra: &addrs.CheckRuleDiagnosticExtra{
+			CheckRule: addr,
+		},
+	})
 
 	return checkResult{
 		Status:         status,

internal/terraform/node_check.go

@@ -173,8 +173,8 @@ func (n *nodeCheckAssert) Execute(ctx EvalContext, _ walkOperation) tfdiags.Diag
 	// Otherwise let's still validate the config and references and return
 	// diagnostics if references do not exist etc.
 	var diags tfdiags.Diagnostics
-	for _, assert := range n.config.Asserts {
-		_, _, moreDiags := validateCheckRule(addrs.CheckAssertion, assert, ctx, n.addr, EvalDataForNoInstanceKey)
+	for ix, assert := range n.config.Asserts {
+		_, _, moreDiags := validateCheckRule(addrs.NewCheckRule(n.addr, addrs.CheckAssertion, ix), assert, ctx, EvalDataForNoInstanceKey)
 		diags = diags.Append(moreDiags)
 	}
 	return diags

internal/terraform/node_resource_abstract_instance.go

@@ -1703,6 +1703,8 @@ func (n *NodeAbstractResourceInstance) planDataSource(ctx EvalContext, checkRule
 	// an apply operation.
 
 	if nested {
+		addr := check.Addr().Absolute(n.Addr.Module)
+
 		// Let's fix things up for a nested data block.
 		//
 		// A nested data block doesn't error, and creates a planned change. So,
@@ -1719,13 +1721,17 @@ func (n *NodeAbstractResourceInstance) planDataSource(ctx EvalContext, checkRule
 
 			// We still want to report the check as failed even if we are still
 			// letting it run again during the apply stage.
-			ctx.Checks().ReportCheckFailure(check.Addr().Absolute(n.Addr.Module), addrs.CheckDataResource, 0, readDiags.Err().Error())
+			ctx.Checks().ReportCheckFailure(addr, addrs.CheckDataResource, 0, readDiags.Err().Error())
 		}
 
 		// Any warning or error diagnostics we'll wrap with some special checks
 		// diagnostics. This is so we can identify them later, and so they'll
 		// only report as warnings.
-		readDiags = tfdiags.AsCheckBlockDiagnostics(readDiags)
+		readDiags = tfdiags.OverrideAll(readDiags, tfdiags.Warning, func() tfdiags.DiagnosticExtraWrapper {
+			return &addrs.CheckRuleDiagnosticExtra{
+				CheckRule: addrs.NewCheckRule(addr, addrs.CheckDataResource, 0),
+			}
+		})
 
 		if !skipPlanChanges {
 			// refreshOnly plans cannot produce planned changes, so we only do
@@ -1882,21 +1888,31 @@ func (n *NodeAbstractResourceInstance) applyDataSource(ctx EvalContext, planned
 
 	newVal, readDiags := n.readDataSource(ctx, configVal)
 	if check, nested := n.nestedInCheckBlock(); nested {
+		addr := check.Addr().Absolute(n.Addr.Module)
+
 		// We're just going to jump in here and hide away any errors for nested
 		// data blocks.
 		if readDiags.HasErrors() {
-			ctx.Checks().ReportCheckFailure(check.Addr().Absolute(n.Addr.Module), addrs.CheckDataResource, 0, readDiags.Err().Error())
-			diags = diags.Append(tfdiags.AsCheckBlockDiagnostics(readDiags))
+			ctx.Checks().ReportCheckFailure(addr, addrs.CheckDataResource, 0, readDiags.Err().Error())
+			diags = diags.Append(tfdiags.OverrideAll(readDiags, tfdiags.Warning, func() tfdiags.DiagnosticExtraWrapper {
+				return &addrs.CheckRuleDiagnosticExtra{
+					CheckRule: addrs.NewCheckRule(addr, addrs.CheckDataResource, 0),
+				}
+			}))
 			return nil, keyData, diags
 		}
 
 		// Even though we know there are no errors here, we still want to
 		// identify these diags has having been generated from a check block.
-		readDiags = tfdiags.AsCheckBlockDiagnostics(readDiags)
+		readDiags = tfdiags.OverrideAll(readDiags, tfdiags.Warning, func() tfdiags.DiagnosticExtraWrapper {
+			return &addrs.CheckRuleDiagnosticExtra{
+				CheckRule: addrs.NewCheckRule(addr, addrs.CheckDataResource, 0),
+			}
+		})
 
 		// If no errors, just remember to report this as a success and continue
 		// as normal.
-		ctx.Checks().ReportCheckResult(check.Addr().Absolute(n.Addr.Module), addrs.CheckDataResource, 0, checks.StatusPass)
+		ctx.Checks().ReportCheckResult(addr, addrs.CheckDataResource, 0, checks.StatusPass)
 	}
 
 	diags = diags.Append(readDiags)

internal/tfdiags/checks.go

@@ -1,76 +0,0 @@
-// Copyright (c) HashiCorp, Inc.
-// SPDX-License-Identifier: MPL-2.0
-
-package tfdiags
-
-import (
-	"fmt"
-
-	"github.com/hashicorp/hcl/v2"
-)
-
-var _ Diagnostic = CheckBlockDiagnostic{}
-
-// CheckBlockDiagnostic is a diagnostic produced by a Terraform config Check block.
-//
-// It only ever returns warnings, and will not be consolidated as part of the
-// Diagnostics.ConsolidateWarnings function.
-type CheckBlockDiagnostic struct {
-	diag Diagnostic
-}
-
-// AsCheckBlockDiagnostics will wrap every diagnostic in diags in a
-// CheckBlockDiagnostic.
-func AsCheckBlockDiagnostics(diags Diagnostics) Diagnostics {
-	if len(diags) == 0 {
-		return nil
-	}
-
-	ret := make(Diagnostics, len(diags))
-	for i, diag := range diags {
-		ret[i] = CheckBlockDiagnostic{diag}
-	}
-	return ret
-}
-
-// AsCheckBlockDiagnostic will wrap a Diagnostic or a hcl.Diagnostic in a
-// CheckBlockDiagnostic.
-func AsCheckBlockDiagnostic(diag interface{}) Diagnostic {
-	switch d := diag.(type) {
-	case Diagnostic:
-		return CheckBlockDiagnostic{d}
-	case *hcl.Diagnostic:
-		return CheckBlockDiagnostic{hclDiagnostic{d}}
-	default:
-		panic(fmt.Errorf("can't construct diagnostic from %T", diag))
-	}
-}
-
-// IsFromCheckBlock returns true if the specified Diagnostic is a
-// CheckBlockDiagnostic.
-func IsFromCheckBlock(diag Diagnostic) bool {
-	_, ok := diag.(CheckBlockDiagnostic)
-	return ok
-}
-
-func (c CheckBlockDiagnostic) Severity() Severity {
-	// Regardless of the severity of the underlying diagnostic, check blocks
-	// only ever report Warning severity.
-	return Warning
-}
-
-func (c CheckBlockDiagnostic) Description() Description {
-	return c.diag.Description()
-}
-
-func (c CheckBlockDiagnostic) Source() Source {
-	return c.diag.Source()
-}
-
-func (c CheckBlockDiagnostic) FromExpr() *FromExpr {
-	return c.diag.FromExpr()
-}
-
-func (c CheckBlockDiagnostic) ExtraInfo() interface{} {
-	return c.diag.ExtraInfo()
-}

internal/tfdiags/consolidate_warnings.go

@@ -46,9 +46,8 @@ func (diags Diagnostics) ConsolidateWarnings(threshold int) Diagnostics {
 			continue
 		}
 
-		if _, ok := diag.(CheckBlockDiagnostic); ok {
-			// Check diagnostics are never consolidated, the users have asked
-			// to be informed about each of these.
+		if DoNotConsolidateDiagnostic(diag) {
+			// Then do not consolidate this diagnostic.
 			newDiags = newDiags.Append(diag)
 			continue
 		}

internal/tfdiags/consolidate_warnings_test.go

@@ -54,6 +54,44 @@ func TestConsolidateWarnings(t *testing.T) {
 		},
 	})
 
+	// Finally, we'll just add a set of diags that should not be consolidated.
+
+	diags = diags.Append(&hcl.Diagnostic{
+		Severity: hcl.DiagWarning,
+		Summary:  "do not consolidate",
+		Detail:   "warning 1, I should not have been consolidated",
+		Subject: &hcl.Range{
+			Filename: "bar.tf",
+			Start:    hcl.Pos{Line: 1, Column: 1, Byte: 0},
+			End:      hcl.Pos{Line: 1, Column: 1, Byte: 0},
+		},
+		Extra: doNotConsolidate(true),
+	})
+
+	diags = diags.Append(&hcl.Diagnostic{
+		Severity: hcl.DiagWarning,
+		Summary:  "do not consolidate",
+		Detail:   "warning 2, I should not have been consolidated",
+		Subject: &hcl.Range{
+			Filename: "bar.tf",
+			Start:    hcl.Pos{Line: 1, Column: 1, Byte: 0},
+			End:      hcl.Pos{Line: 1, Column: 1, Byte: 0},
+		},
+		Extra: doNotConsolidate(true),
+	})
+
+	diags = diags.Append(&hcl.Diagnostic{
+		Severity: hcl.DiagWarning,
+		Summary:  "do not consolidate",
+		Detail:   "warning 3, I should not have been consolidated",
+		Subject: &hcl.Range{
+			Filename: "bar.tf",
+			Start:    hcl.Pos{Line: 1, Column: 1, Byte: 0},
+			End:      hcl.Pos{Line: 1, Column: 1, Byte: 0},
+		},
+		Extra: doNotConsolidate(true),
+	})
+
 	// We're using ForRPC here to force the diagnostics to be of a consistent
 	// type that we can easily assert against below.
 	got := diags.ConsolidateWarnings(2).ForRPC()
@@ -174,9 +212,50 @@ func TestConsolidateWarnings(t *testing.T) {
 				End:      SourcePos{Line: 1, Column: 1, Byte: 0},
 			},
 		},
+
+		// The final set of warnings should not have been consolidated because
+		// of our filter function.
+		&rpcFriendlyDiag{
+			Severity_: Warning,
+			Summary_:  "do not consolidate",
+			Detail_:   "warning 1, I should not have been consolidated",
+			Subject_: &SourceRange{
+				Filename: "bar.tf",
+				Start:    SourcePos{Line: 1, Column: 1, Byte: 0},
+				End:      SourcePos{Line: 1, Column: 1, Byte: 0},
+			},
+		},
+		&rpcFriendlyDiag{
+			Severity_: Warning,
+			Summary_:  "do not consolidate",
+			Detail_:   "warning 2, I should not have been consolidated",
+			Subject_: &SourceRange{
+				Filename: "bar.tf",
+				Start:    SourcePos{Line: 1, Column: 1, Byte: 0},
+				End:      SourcePos{Line: 1, Column: 1, Byte: 0},
+			},
+		},
+		&rpcFriendlyDiag{
+			Severity_: Warning,
+			Summary_:  "do not consolidate",
+			Detail_:   "warning 3, I should not have been consolidated",
+			Subject_: &SourceRange{
+				Filename: "bar.tf",
+				Start:    SourcePos{Line: 1, Column: 1, Byte: 0},
+				End:      SourcePos{Line: 1, Column: 1, Byte: 0},
+			},
+		},
 	}
 
 	if diff := cmp.Diff(want, got); diff != "" {
 		t.Errorf("wrong result\n%s", diff)
 	}
 }
+
+type doNotConsolidate bool
+
+var _ DiagnosticExtraDoNotConsolidate = doNotConsolidate(true)
+
+func (d doNotConsolidate) DoNotConsolidateDiagnostic() bool {
+	return bool(d)
+}

internal/tfdiags/diagnostic_extra.go

@@ -105,6 +105,14 @@ type DiagnosticExtraUnwrapper interface {
 	UnwrapDiagnosticExtra() interface{}
 }
 
+// DiagnosticExtraWrapper is an interface implemented by values that can be
+// dynamically updated to wrap other extra info.
+type DiagnosticExtraWrapper interface {
+	// WrapDiagnosticExtra accepts an ExtraInfo that it should add within the
+	// current ExtraInfo.
+	WrapDiagnosticExtra(inner interface{})
+}
+
 // DiagnosticExtraBecauseUnknown is an interface implemented by values in
 // the Extra field of Diagnostic when the diagnostic is potentially caused by
 // the presence of unknown values in an expression evaluation.
@@ -172,3 +180,22 @@ func DiagnosticCausedBySensitive(diag Diagnostic) bool {
 	}
 	return maybe.DiagnosticCausedBySensitive()
 }
+
+// DiagnosticExtraDoNotConsolidate tells the Diagnostics.ConsolidateWarnings
+// function not to consolidate this diagnostic if it otherwise would.
+type DiagnosticExtraDoNotConsolidate interface {
+	// DoNotConsolidateDiagnostic returns true if the associated diagnostic
+	// should not be consolidated by the Diagnostics.ConsolidateWarnings
+	// function.
+	DoNotConsolidateDiagnostic() bool
+}
+
+// DoNotConsolidateDiagnostic returns true if the given diagnostic should not
+// be consolidated by the Diagnostics.ConsolidateWarnings function.
+func DoNotConsolidateDiagnostic(diag Diagnostic) bool {
+	maybe := ExtraInfo[DiagnosticExtraDoNotConsolidate](diag)
+	if maybe == nil {
+		return false
+	}
+	return maybe.DoNotConsolidateDiagnostic()
+}

internal/tfdiags/override.go

@@ -0,0 +1,72 @@
+package tfdiags
+
+// overriddenDiagnostic implements the Diagnostic interface by wrapping another
+// Diagnostic while overriding the severity of the original Diagnostic.
+type overriddenDiagnostic struct {
+	original Diagnostic
+	severity Severity
+	extra    interface{}
+}
+
+var _ Diagnostic = overriddenDiagnostic{}
+
+// OverrideAll accepts a set of Diagnostics and wraps them with a new severity
+// and, optionally, a new ExtraInfo.
+func OverrideAll(originals Diagnostics, severity Severity, createExtra func() DiagnosticExtraWrapper) Diagnostics {
+	var diags Diagnostics
+	for _, diag := range originals {
+		diags = diags.Append(Override(diag, severity, createExtra))
+	}
+	return diags
+}
+
+// Override matches OverrideAll except it operates over a single Diagnostic
+// rather than multiple Diagnostics.
+func Override(original Diagnostic, severity Severity, createExtra func() DiagnosticExtraWrapper) Diagnostic {
+	extra := original.ExtraInfo()
+	if createExtra != nil {
+		nw := createExtra()
+		nw.WrapDiagnosticExtra(extra)
+		extra = nw
+	}
+
+	return overriddenDiagnostic{
+		original: original,
+		severity: severity,
+		extra:    extra,
+	}
+}
+
+// UndoOverride will return the original diagnostic that was overridden within
+// the OverrideAll function.
+//
+// If the provided Diagnostic was never overridden then it is simply returned
+// unchanged.
+func UndoOverride(diag Diagnostic) Diagnostic {
+	if override, ok := diag.(overriddenDiagnostic); ok {
+		return override.original
+	}
+
+	// Then it wasn't overridden, so we'll just return the diag unchanged.
+	return diag
+}
+
+func (o overriddenDiagnostic) Severity() Severity {
+	return o.severity
+}
+
+func (o overriddenDiagnostic) Description() Description {
+	return o.original.Description()
+}
+
+func (o overriddenDiagnostic) Source() Source {
+	return o.original.Source()
+}
+
+func (o overriddenDiagnostic) FromExpr() *FromExpr {
+	return o.original.FromExpr()
+}
+
+func (o overriddenDiagnostic) ExtraInfo() interface{} {
+	return o.extra
+}

internal/tfdiags/override_test.go

@@ -0,0 +1,80 @@
+package tfdiags
+
+import (
+	"testing"
+
+	"github.com/hashicorp/hcl/v2"
+)
+
+func TestOverride_UpdatesSeverity(t *testing.T) {
+	original := Sourceless(Error, "summary", "detail")
+	override := Override(original, Warning, nil)
+
+	if override.Severity() != Warning {
+		t.Errorf("expected warning but was %s", override.Severity())
+	}
+}
+
+func TestOverride_MaintainsExtra(t *testing.T) {
+	original := hclDiagnostic{&hcl.Diagnostic{
+		Severity: hcl.DiagError,
+		Summary:  "summary",
+		Detail:   "detail",
+		Extra:    "extra",
+	}}
+	override := Override(original, Warning, nil)
+
+	if override.ExtraInfo().(string) != "extra" {
+		t.Errorf("invalid extra info %v", override.ExtraInfo())
+	}
+}
+
+func TestOverride_WrapsExtra(t *testing.T) {
+	original := hclDiagnostic{&hcl.Diagnostic{
+		Severity: hcl.DiagError,
+		Summary:  "summary",
+		Detail:   "detail",
+		Extra:    "extra",
+	}}
+	override := Override(original, Warning, func() DiagnosticExtraWrapper {
+		return &extraWrapper{
+			mine: "mine",
+		}
+	})
+
+	wrapper := override.ExtraInfo().(*extraWrapper)
+	if wrapper.mine != "mine" {
+		t.Errorf("invalid extra info %v", override.ExtraInfo())
+	}
+	if wrapper.original.(string) != "extra" {
+		t.Errorf("invalid wrapped extra info %v", override.ExtraInfo())
+	}
+}
+
+func TestUndoOverride(t *testing.T) {
+	original := Sourceless(Error, "summary", "detail")
+	override := Override(original, Warning, nil)
+	restored := UndoOverride(override)
+
+	if restored.Severity() != Error {
+		t.Errorf("expected warning but was %s", restored.Severity())
+	}
+}
+
+func TestUndoOverride_NotOverridden(t *testing.T) {
+	original := Sourceless(Error, "summary", "detail")
+	restored := UndoOverride(original) // Shouldn't do anything bad.
+
+	if restored.Severity() != Error {
+		t.Errorf("expected warning but was %s", restored.Severity())
+	}
+}
+
+type extraWrapper struct {
+	mine     string
+	original interface{}
+}
+
+func (e *extraWrapper) WrapDiagnosticExtra(inner interface{}) {
+	e.original = inner
+}