From 083abaac098cc168f7624e905569914ed5b277ab Mon Sep 17 00:00:00 2001
From: Eugene Burkov <e.burkov@adguard.com>
Date: Wed, 6 Dec 2023 18:56:14 +0300
Subject: [PATCH] Pull request 2099: improve validator test

Squashed commit of the following:

commit 3a7b8fd5e73cd301643d3a907f38b2dade87b60d
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Wed Dec 6 18:36:08 2023 +0300

    dnsforward: imp code

commit 9751bf513ddf7013071e22500a347cdda1540ce6
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Wed Dec 6 18:27:59 2023 +0300

    dnsforward: imp validator test
---
 .../dnsforward/upstreams_internal_test.go     | 36 +++++++++++++------
 1 file changed, 25 insertions(+), 11 deletions(-)

diff --git a/internal/dnsforward/upstreams_internal_test.go b/internal/dnsforward/upstreams_internal_test.go
index d3d4ebb6..9bf2f5cf 100644
--- a/internal/dnsforward/upstreams_internal_test.go
+++ b/internal/dnsforward/upstreams_internal_test.go
@@ -4,7 +4,6 @@ import (
 	"net"
 	"net/url"
 	"strings"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -164,13 +163,16 @@ func TestUpstreamConfigValidator(t *testing.T) {
 }
 
 func TestUpstreamConfigValidator_Check_once(t *testing.T) {
-	reqs := atomic.Int32{}
-	reset := func() { reqs.Store(0) }
+	type signal = struct{}
 
+	reqCh := make(chan signal)
 	hdlr := dns.HandlerFunc(func(w dns.ResponseWriter, m *dns.Msg) {
+		pt := testutil.PanicT{}
+
 		err := w.WriteMsg(new(dns.Msg).SetReply(m))
-		require.NoError(testutil.PanicT{}, err)
-		reqs.Add(1)
+		require.NoError(pt, err)
+
+		testutil.RequireSend(pt, reqCh, signal{}, testTimeout)
 	})
 
 	addr := (&url.URL{
@@ -179,6 +181,10 @@ func TestUpstreamConfigValidator_Check_once(t *testing.T) {
 	}).String()
 	twoAddrs := strings.Join([]string{addr, addr}, " ")
 
+	wantStatus := map[string]string{
+		addr: "OK",
+	}
+
 	testCases := []struct {
 		name string
 		ups  []string
@@ -195,15 +201,23 @@ func TestUpstreamConfigValidator_Check_once(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			t.Cleanup(reset)
-
 			cv := newUpstreamConfigValidator(tc.ups, nil, nil, &upstream.Options{
-				Timeout: 100 * time.Millisecond,
+				Timeout: testTimeout,
 			})
-			cv.check()
-			cv.close()
 
-			assert.Equal(t, int32(1), reqs.Load())
+			go func() {
+				cv.check()
+				testutil.RequireSend(testutil.PanicT{}, reqCh, signal{}, testTimeout)
+			}()
+
+			// Wait for the only request to be sent.
+			testutil.RequireReceive(t, reqCh, testTimeout)
+
+			// Wait for the check to finish.
+			testutil.RequireReceive(t, reqCh, testTimeout)
+
+			cv.close()
+			require.Equal(t, wantStatus, cv.status())
 		})
 	}
 }