Skips data sources from both original state and imported state

helper/resource/testing_new_import_state.go

@@ -158,20 +158,27 @@ func testStepNewImportState(ctx context.Context, t testing.T, helper *plugintest
 	if step.ImportStateVerify {
 		logging.HelperResourceTrace(ctx, "Using TestStep ImportStateVerify")
 
-		newResources := importState.RootModule().Resources
-		oldResources := state.RootModule().Resources
+		// Ensure that we do not match against data sources as they
+		// cannot be imported and are not what we want to verify.
+		// Mode is not present in ResourceState so we use the
+		// stringified ResourceStateKey for comparison.
+		newResources := make(map[string]*terraform.ResourceState)
+		for k, v := range importState.RootModule().Resources {
+			if !strings.HasPrefix(k, "data.") {
+				newResources[k] = v
+			}
+		}
+		oldResources := make(map[string]*terraform.ResourceState)
+		for k, v := range state.RootModule().Resources {
+			if !strings.HasPrefix(k, "data.") {
+				oldResources[k] = v
+			}
+		}
 
 		for _, r := range newResources {
 			// Find the existing resource
 			var oldR *terraform.ResourceState
-			for r2Key, r2 := range oldResources {
-				// Ensure that we do not match against data sources as they
-				// cannot be imported and are not what we want to verify.
-				// Mode is not present in ResourceState so we use the
-				// stringified ResourceStateKey for comparison.
-				if strings.HasPrefix(r2Key, "data.") {
-					continue
-				}
+			for _, r2 := range oldResources {
 
 				if r2.Primary != nil && r2.Primary.ID == r.Primary.ID && r2.Type == r.Type && r2.Provider == r.Provider {
 					oldR = r2