From e662b4f862edd4a60cb32c7dd95f0059a3253c2d Mon Sep 17 00:00:00 2001
From: Julian Tibble <tibbes@github.com>
Date: Sun, 28 Feb 2021 22:52:18 +0000
Subject: [PATCH] Add test for importing github_branch_protection

There was an existing test for importing by <repo-name>:<pattern>,
but not a test for <repo-node-id>:<pattern>.
---
 .../resource_github_branch_protection_test.go | 28 +++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/github/resource_github_branch_protection_test.go b/github/resource_github_branch_protection_test.go
index f6cbc1ed29..e56d2eadbc 100644
--- a/github/resource_github_branch_protection_test.go
+++ b/github/resource_github_branch_protection_test.go
@@ -62,7 +62,7 @@ func TestAccGithubBranchProtection(t *testing.T) {
 						ResourceName:      "github_branch_protection.test",
 						ImportState:       true,
 						ImportStateVerify: true,
-						ImportStateIdFunc: branchProtectionImportStateIdFunc(
+						ImportStateIdFunc: importBranchProtectionByRepoName(
 							fmt.Sprintf("tf-acc-test-%s", randomID), "main",
 						),
 					},
@@ -122,6 +122,13 @@ func TestAccGithubBranchProtection(t *testing.T) {
 						Config: config,
 						Check:  check,
 					},
+					{
+						ResourceName:      "github_branch_protection.test",
+						ImportState:       true,
+						ImportStateVerify: true,
+						ImportStateIdFunc: importBranchProtectionByRepoID(
+							"github_repository.test", "main"),
+					},
 				},
 			})
 		}
@@ -322,8 +329,25 @@ func TestAccGithubBranchProtection(t *testing.T) {
 
 }
 
-func branchProtectionImportStateIdFunc(repo, pattern string) resource.ImportStateIdFunc {
+func importBranchProtectionByRepoName(repo, pattern string) resource.ImportStateIdFunc {
+	// test importing using an ID of the form <repo-name>:<branch-protection-pattern>
 	return func(s *terraform.State) (string, error) {
 		return fmt.Sprintf("%s:%s", repo, pattern), nil
 	}
 }
+
+func importBranchProtectionByRepoID(repoLogicalName, pattern string) resource.ImportStateIdFunc {
+	// test importing using an ID of the form <repo-node-id>:<branch-protection-pattern>
+	// by retrieving the GraphQL ID from the terraform.State
+	return func(s *terraform.State) (string, error) {
+		repo := s.RootModule().Resources[repoLogicalName]
+		if repo == nil {
+			return "", fmt.Errorf("Cannot find %s in terraform state", repoLogicalName)
+		}
+		repoID, found := repo.Primary.Attributes["node_id"]
+		if !found {
+			return "", fmt.Errorf("Repository %s does not have a node_id in terraform state", repo.Primary.ID)
+		}
+		return fmt.Sprintf("%s:%s", repoID, pattern), nil
+	}
+}