From 6273f9b38e22632b3b36b3cd142ad461f97c5e96 Mon Sep 17 00:00:00 2001
From: Terry Howe <terrylhowe@gmail.com>
Date: Wed, 27 Aug 2025 08:18:26 -0600
Subject: [PATCH] fix: flaky registry data race on mockdns close

Signed-off-by: Terry Howe <terrylhowe@gmail.com>
---
 pkg/registry/client_http_test.go         |  5 +----
 pkg/registry/client_insecure_tls_test.go |  5 +----
 pkg/registry/client_tls_test.go          |  5 +----
 pkg/registry/utils_test.go               | 20 ++++++++++----------
 4 files changed, 13 insertions(+), 22 deletions(-)

diff --git a/pkg/registry/client_http_test.go b/pkg/registry/client_http_test.go
index 043fd4205..dddd29ee9 100644
--- a/pkg/registry/client_http_test.go
+++ b/pkg/registry/client_http_test.go
@@ -32,10 +32,7 @@ type HTTPRegistryClientTestSuite struct {
 
 func (suite *HTTPRegistryClientTestSuite) SetupSuite() {
 	// init test client
-	dockerRegistry := setup(&suite.TestSuite, false, false)
-
-	// Start Docker registry
-	go dockerRegistry.ListenAndServe()
+	setup(&suite.TestSuite, false, false)
 }
 
 func (suite *HTTPRegistryClientTestSuite) TearDownSuite() {
diff --git a/pkg/registry/client_insecure_tls_test.go b/pkg/registry/client_insecure_tls_test.go
index accbf1670..03354475a 100644
--- a/pkg/registry/client_insecure_tls_test.go
+++ b/pkg/registry/client_insecure_tls_test.go
@@ -29,10 +29,7 @@ type InsecureTLSRegistryClientTestSuite struct {
 
 func (suite *InsecureTLSRegistryClientTestSuite) SetupSuite() {
 	// init test client
-	dockerRegistry := setup(&suite.TestSuite, true, true)
-
-	// Start Docker registry
-	go dockerRegistry.ListenAndServe()
+	setup(&suite.TestSuite, true, true)
 }
 
 func (suite *InsecureTLSRegistryClientTestSuite) TearDownSuite() {
diff --git a/pkg/registry/client_tls_test.go b/pkg/registry/client_tls_test.go
index 0897858b5..2bf1750a9 100644
--- a/pkg/registry/client_tls_test.go
+++ b/pkg/registry/client_tls_test.go
@@ -31,10 +31,7 @@ type TLSRegistryClientTestSuite struct {
 
 func (suite *TLSRegistryClientTestSuite) SetupSuite() {
 	// init test client
-	dockerRegistry := setup(&suite.TestSuite, true, false)
-
-	// Start Docker registry
-	go dockerRegistry.ListenAndServe()
+	setup(&suite.TestSuite, true, false)
 }
 
 func (suite *TLSRegistryClientTestSuite) TearDownSuite() {
diff --git a/pkg/registry/utils_test.go b/pkg/registry/utils_test.go
index b46317fc6..781f3dd75 100644
--- a/pkg/registry/utils_test.go
+++ b/pkg/registry/utils_test.go
@@ -29,7 +29,6 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/distribution/distribution/v3/configuration"
@@ -65,12 +64,13 @@ type TestSuite struct {
 	CompromisedRegistryHost string
 	WorkspaceDir            string
 	RegistryClient          *Client
+	dockerRegistry          *registry.Registry
 
 	// A mock DNS server needed for TLS connection testing.
 	srv *mockdns.Server
 }
 
-func setup(suite *TestSuite, tlsEnabled, insecure bool) *registry.Registry {
+func setup(suite *TestSuite, tlsEnabled, insecure bool) {
 	suite.WorkspaceDir = testWorkspaceDir
 	os.RemoveAll(suite.WorkspaceDir)
 	os.Mkdir(suite.WorkspaceDir, 0700)
@@ -166,20 +166,20 @@ func setup(suite *TestSuite, tlsEnabled, insecure bool) *registry.Registry {
 			config.HTTP.TLS.ClientCAs = []string{tlsCA}
 		}
 	}
-	dockerRegistry, err := registry.NewRegistry(context.Background(), config)
+	suite.dockerRegistry, err = registry.NewRegistry(context.Background(), config)
 	suite.Nil(err, "no error creating test registry")
 
 	suite.CompromisedRegistryHost = initCompromisedRegistryTestServer()
-	return dockerRegistry
+	go func() {
+		_ = suite.dockerRegistry.ListenAndServe()
+		_ = suite.srv.Close()
+		mockdns.UnpatchNet(net.DefaultResolver)
+	}()
 }
 
 func teardown(suite *TestSuite) {
-	var lock sync.Mutex
-	lock.Lock()
-	defer lock.Unlock()
-	if suite.srv != nil {
-		mockdns.UnpatchNet(net.DefaultResolver)
-		suite.srv.Close()
+	if suite.dockerRegistry != nil {
+		_ = suite.dockerRegistry.Shutdown(context.Background())
 	}
 }