From 52f121772f633909447dd44f30984cd4f189ef96 Mon Sep 17 00:00:00 2001
From: Arran Schlosberg <me@arranschlosberg.com>
Date: Wed, 24 Jul 2024 20:13:52 +0100
Subject: [PATCH] chore: allow test-only imports in `*test` and `/tests/**`
 packages

---
 scripts/lint.sh                               | 15 +++++++++++----
 snow/consensus/snowman/snowmantest/block.go   |  2 --
 snow/consensus/snowman/snowmantest/engine.go  |  2 --
 snow/consensus/snowman/snowmantest/require.go |  2 --
 snow/snowtest/context.go                      |  2 --
 snow/snowtest/decidable.go                    |  2 --
 snow/snowtest/status.go                       |  2 --
 7 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/scripts/lint.sh b/scripts/lint.sh
index 74c5f47555f..123ba982f0a 100755
--- a/scripts/lint.sh
+++ b/scripts/lint.sh
@@ -88,25 +88,32 @@ function test_interface_compliance_nil {
 
 function test_import_testing_only_in_tests {
   ROOT=$( git rev-parse --show-toplevel )
-  NON_TEST_GO_FILES=$( find "${ROOT}" -iname '*.go' ! -iname '*_test.go');
+  NON_TEST_GO_FILES=$( find "${ROOT}" -iname '*.go' ! -iname '*_test.go' ! -path "${ROOT}/tests/*" );
 
   IMPORT_TESTING=$( echo "${NON_TEST_GO_FILES}" | xargs grep -lP '^\s*(import\s+)?"testing"');
   IMPORT_TESTIFY=$( echo "${NON_TEST_GO_FILES}" | xargs grep -l '"github.com/stretchr/testify');
+  IMPORT_FROM_TESTS=$( echo "${NON_TEST_GO_FILES}" | xargs grep -l '"github.com/ava-labs/tests/');
+  IMPORT_TEST_PKG=$( echo "${NON_TEST_GO_FILES}" | xargs grep -lP '"github.com/ava-labs/.*?test"');
+
   # TODO(arr4n): send a PR to add support for build tags in `mockgen` and then enable this.
   # IMPORT_GOMOCK=$( echo "${NON_TEST_GO_FILES}" | xargs grep -l '"go.uber.org/mock');
-  HAVE_TEST_LOGIC=$( printf "%s\n%s" "${IMPORT_TESTING}" "${IMPORT_TESTIFY}" );
+  HAVE_TEST_LOGIC=$( printf "%s\n%s\n%s\n%s" "${IMPORT_TESTING}" "${IMPORT_TESTIFY}" "${IMPORT_FROM_TESTS}" "${IMPORT_TEST_PKG}" );
 
   TAGGED_AS_TEST=$( echo "${NON_TEST_GO_FILES}" | xargs grep -lP '^\/\/go:build\s+(.+(,|\s+))?test[,\s]?');
+  IN_TEST_PKG=$( echo "${NON_TEST_GO_FILES}" | grep -P '.*test/[^/]+\.go$' ) # directory (hence package name) ends in "test"
+
+  # Files in /tests/ are already excluded by the `find ... ! -path`
+  INTENDED_FOR_TESTING=$( printf "%s\n%s" "${TAGGED_AS_TEST}" "${IN_TEST_PKG}" )
 
   # -3 suppresses files that have test logic and have the "test" build tag
   # -2 suppresses files that are tagged despite not having detectable test logic
-  UNTAGGED=$( comm -23 <( echo "${HAVE_TEST_LOGIC}" | sort -u ) <( echo "${TAGGED_AS_TEST}" | sort -u ) );
+  UNTAGGED=$( comm -23 <( echo "${HAVE_TEST_LOGIC}" | sort -u ) <( echo "${INTENDED_FOR_TESTING}" | sort -u ) );
   if [ -z "${UNTAGGED}" ];
   then
     return 0;
   fi
 
-  echo "Non-test Go files importing test-only packages MUST have '//go:build test' tag:";
+  echo 'Non-test Go files importing test-only packages MUST (a) have '//go:build test' tag; (b) be in *test package; or (c) be in /tests/ directory:';
   echo "${UNTAGGED}";
   return 1;
 }
diff --git a/snow/consensus/snowman/snowmantest/block.go b/snow/consensus/snowman/snowmantest/block.go
index 19348532e54..4dba3d0cc69 100644
--- a/snow/consensus/snowman/snowmantest/block.go
+++ b/snow/consensus/snowman/snowmantest/block.go
@@ -1,8 +1,6 @@
 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-//go:build test
-
 package snowmantest
 
 import (
diff --git a/snow/consensus/snowman/snowmantest/engine.go b/snow/consensus/snowman/snowmantest/engine.go
index 954b6bbe707..ee7560eadc5 100644
--- a/snow/consensus/snowman/snowmantest/engine.go
+++ b/snow/consensus/snowman/snowmantest/engine.go
@@ -1,8 +1,6 @@
 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-//go:build test
-
 package snowmantest
 
 import (
diff --git a/snow/consensus/snowman/snowmantest/require.go b/snow/consensus/snowman/snowmantest/require.go
index 5a5e5d09bf4..7b942740dfd 100644
--- a/snow/consensus/snowman/snowmantest/require.go
+++ b/snow/consensus/snowman/snowmantest/require.go
@@ -1,8 +1,6 @@
 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-//go:build test
-
 package snowmantest
 
 import (
diff --git a/snow/snowtest/context.go b/snow/snowtest/context.go
index be08ae75b21..3cacc8e873b 100644
--- a/snow/snowtest/context.go
+++ b/snow/snowtest/context.go
@@ -1,8 +1,6 @@
 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-//go:build test
-
 package snowtest
 
 import (
diff --git a/snow/snowtest/decidable.go b/snow/snowtest/decidable.go
index 3cbdb2c1c9c..691cc90120b 100644
--- a/snow/snowtest/decidable.go
+++ b/snow/snowtest/decidable.go
@@ -1,8 +1,6 @@
 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-//go:build test
-
 package snowtest
 
 import (
diff --git a/snow/snowtest/status.go b/snow/snowtest/status.go
index 36c26f83cbd..6fd90f856a3 100644
--- a/snow/snowtest/status.go
+++ b/snow/snowtest/status.go
@@ -1,8 +1,6 @@
 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-//go:build test
-
 package snowtest
 
 type Status int