resolver/delegatingresolver: wait for proxy resolver build before update in tests (#8304)

Author: eshitachandwani

internal/resolver/delegatingresolver/delegatingresolver_ext_test.go

@@ -131,15 +131,29 @@ func (s) TestDelegatingResolverNoProxyEnvVarsSet(t *testing.T) {
 // overwriting the previously registered DNS resolver. This allows the test to
 // mock the DNS resolution for the proxy resolver. It also registers the
 // original DNS resolver after the test is done.
-func setupDNS(t *testing.T) *manual.Resolver {
+func setupDNS(t *testing.T) (*manual.Resolver, chan struct{}) {
 	t.Helper()
 	mr := manual.NewBuilderWithScheme("dns")
 
 	dnsResolverBuilder := resolver.Get("dns")
 	resolver.Register(mr)
 
+	resolverBuilt := make(chan struct{})
+	mr.BuildCallback = func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) {
+		close(resolverBuilt)
+	}
+
 	t.Cleanup(func() { resolver.Register(dnsResolverBuilder) })
-	return mr
+	return mr, resolverBuilt
+}
+
+func mustBuildResolver(ctx context.Context, t *testing.T, buildCh chan struct{}) {
+	t.Helper()
+	select {
+	case <-buildCh:
+	case <-ctx.Done():
+		t.Fatalf("Context timed out waiting for resolver to be built.")
+	}
 }
 
 // proxyAddressWithTargetAttribute creates a resolver.Address for the proxy,
@@ -181,7 +195,7 @@ func (s) TestDelegatingResolverwithDNSAndProxyWithTargetResolution(t *testing.T)
 	targetResolver := manual.NewBuilderWithScheme("dns")
 	target := targetResolver.Scheme() + ":///" + targetTestAddr
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
+	proxyResolver, proxyResolverBuilt := setupDNS(t)
 
 	tcc, stateCh, _ := createTestResolverClientConn(t)
 	if _, err := delegatingresolver.New(resolver.Target{URL: *testutils.MustParseURL(target)}, tcc, resolver.BuildOptions{}, targetResolver, true); err != nil {
@@ -202,6 +216,11 @@ func (s) TestDelegatingResolverwithDNSAndProxyWithTargetResolution(t *testing.T)
 	case <-time.After(defaultTestShortTimeout):
 	}
 
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Wait for the proxy resolver to be built before calling UpdateState.
+	mustBuildResolver(ctx, t, proxyResolverBuilt)
 	proxyResolver.UpdateState(resolver.State{
 		Addresses:     []resolver.Address{{Addr: resolvedProxyTestAddr1}},
 		ServiceConfig: &serviceconfig.ParseResult{},
@@ -218,8 +237,8 @@ func (s) TestDelegatingResolverwithDNSAndProxyWithTargetResolution(t *testing.T)
 	var gotState resolver.State
 	select {
 	case gotState = <-stateCh:
-	case <-time.After(defaultTestTimeout):
-		t.Fatal("Timeout when waiting for a state update from the delegating resolver")
+	case <-ctx.Done():
+		t.Fatal("Context timeed out when waiting for a state update from the delegating resolver")
 	}
 
 	if diff := cmp.Diff(gotState, wantState); diff != "" {
@@ -257,13 +276,18 @@ func (s) TestDelegatingResolverwithDNSAndProxyWithNoTargetResolution(t *testing.
 	targetResolver := manual.NewBuilderWithScheme("dns")
 	target := targetResolver.Scheme() + ":///" + targetTestAddr
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
+	proxyResolver, proxyResolverBuilt := setupDNS(t)
 
 	tcc, stateCh, _ := createTestResolverClientConn(t)
 	if _, err := delegatingresolver.New(resolver.Target{URL: *testutils.MustParseURL(target)}, tcc, resolver.BuildOptions{}, targetResolver, false); err != nil {
 		t.Fatalf("Failed to create delegating resolver: %v", err)
 	}
 
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Wait for the proxy resolver to be built before calling UpdateState.
+	mustBuildResolver(ctx, t, proxyResolverBuilt)
 	proxyResolver.UpdateState(resolver.State{
 		Addresses: []resolver.Address{
 			{Addr: resolvedProxyTestAddr1},
@@ -278,8 +302,8 @@ func (s) TestDelegatingResolverwithDNSAndProxyWithNoTargetResolution(t *testing.
 	var gotState resolver.State
 	select {
 	case gotState = <-stateCh:
-	case <-time.After(defaultTestTimeout):
-		t.Fatal("Timeout when waiting for a state update from the delegating resolver")
+	case <-ctx.Done():
+		t.Fatal("Context timed out when waiting for a state update from the delegating resolver")
 	}
 
 	if diff := cmp.Diff(gotState, wantState); diff != "" {
@@ -319,7 +343,7 @@ func (s) TestDelegatingResolverwithCustomResolverAndProxy(t *testing.T) {
 	targetResolver := manual.NewBuilderWithScheme("test")
 	target := targetResolver.Scheme() + ":///" + targetTestAddr
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
+	proxyResolver, proxyResolverBuilt := setupDNS(t)
 
 	tcc, stateCh, _ := createTestResolverClientConn(t)
 	if _, err := delegatingresolver.New(resolver.Target{URL: *testutils.MustParseURL(target)}, tcc, resolver.BuildOptions{}, targetResolver, false); err != nil {
@@ -340,6 +364,11 @@ func (s) TestDelegatingResolverwithCustomResolverAndProxy(t *testing.T) {
 	case <-time.After(defaultTestShortTimeout):
 	}
 
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Wait for the proxy resolver to be built before calling UpdateState.
+	mustBuildResolver(ctx, t, proxyResolverBuilt)
 	proxyResolver.UpdateState(resolver.State{
 		Addresses:     []resolver.Address{{Addr: resolvedProxyTestAddr1}},
 		ServiceConfig: &serviceconfig.ParseResult{},
@@ -355,8 +384,8 @@ func (s) TestDelegatingResolverwithCustomResolverAndProxy(t *testing.T) {
 	var gotState resolver.State
 	select {
 	case gotState = <-stateCh:
-	case <-time.After(defaultTestTimeout):
-		t.Fatal("Timeout when waiting for a state update from the delegating resolver")
+	case <-ctx.Done():
+		t.Fatal("Context timed out when waiting for a state update from the delegating resolver")
 	}
 
 	if diff := cmp.Diff(gotState, wantState); diff != "" {
@@ -401,7 +430,7 @@ func (s) TestDelegatingResolverForEndpointsWithProxy(t *testing.T) {
 	targetResolver := manual.NewBuilderWithScheme("test")
 	target := targetResolver.Scheme() + ":///" + targetTestAddr
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
+	proxyResolver, proxyResolverBuilt := setupDNS(t)
 
 	tcc, stateCh, _ := createTestResolverClientConn(t)
 	if _, err := delegatingresolver.New(resolver.Target{URL: *testutils.MustParseURL(target)}, tcc, resolver.BuildOptions{}, targetResolver, false); err != nil {
@@ -429,6 +458,11 @@ func (s) TestDelegatingResolverForEndpointsWithProxy(t *testing.T) {
 	case <-time.After(defaultTestShortTimeout):
 	}
 
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Wait for the proxy resolver to be built before calling UpdateState.
+	mustBuildResolver(ctx, t, proxyResolverBuilt)
 	proxyResolver.UpdateState(resolver.State{
 		Endpoints: []resolver.Endpoint{
 			{Addresses: []resolver.Address{{Addr: resolvedProxyTestAddr1}}},
@@ -460,8 +494,8 @@ func (s) TestDelegatingResolverForEndpointsWithProxy(t *testing.T) {
 	var gotState resolver.State
 	select {
 	case gotState = <-stateCh:
-	case <-time.After(defaultTestTimeout):
-		t.Fatal("Timeout when waiting for a state update from the delegating resolver")
+	case <-ctx.Done():
+		t.Fatal("Contex timed out when waiting for a state update from the delegating resolver")
 	}
 
 	if diff := cmp.Diff(gotState, wantState); diff != "" {
@@ -503,8 +537,7 @@ func (s) TestDelegatingResolverForMultipleProxyAddress(t *testing.T) {
 	targetResolver := manual.NewBuilderWithScheme("test")
 	target := targetResolver.Scheme() + ":///" + targetTestAddr
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
-
+	proxyResolver, proxyResolverBuilt := setupDNS(t)
 	tcc, stateCh, _ := createTestResolverClientConn(t)
 	if _, err := delegatingresolver.New(resolver.Target{URL: *testutils.MustParseURL(target)}, tcc, resolver.BuildOptions{}, targetResolver, false); err != nil {
 		t.Fatalf("Failed to create delegating resolver: %v", err)
@@ -524,6 +557,11 @@ func (s) TestDelegatingResolverForMultipleProxyAddress(t *testing.T) {
 	case <-time.After(defaultTestShortTimeout):
 	}
 
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Wait for the proxy resolver to be built before calling UpdateState.
+	mustBuildResolver(ctx, t, proxyResolverBuilt)
 	proxyResolver.UpdateState(resolver.State{
 		Addresses: []resolver.Address{
 			{Addr: resolvedProxyTestAddr1},
@@ -542,8 +580,8 @@ func (s) TestDelegatingResolverForMultipleProxyAddress(t *testing.T) {
 	var gotState resolver.State
 	select {
 	case gotState = <-stateCh:
-	case <-time.After(defaultTestTimeout):
-		t.Fatal("Timeout when waiting for a state update from the delegating resolver")
+	case <-ctx.Done():
+		t.Fatal("Context timed out when waiting for a state update from the delegating resolver")
 	}
 
 	if diff := cmp.Diff(gotState, wantState); diff != "" {
@@ -586,12 +624,8 @@ func (s) TestDelegatingResolverUpdateStateDuringClose(t *testing.T) {
 
 	target := targetResolver.Scheme() + ":///" + "ignored"
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
+	proxyResolver, proxyResolverBuilt := setupDNS(t)
 
-	proxyResolverBuilt := make(chan struct{})
-	proxyResolver.BuildCallback = func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) {
-		close(proxyResolverBuilt)
-	}
 	unblockProxyResolverClose := make(chan struct{}, 1)
 	proxyResolver.CloseCallback = func() {
 		<-unblockProxyResolverClose
@@ -614,11 +648,7 @@ func (s) TestDelegatingResolverUpdateStateDuringClose(t *testing.T) {
 	// Wait for the proxy resolver to be built before calling Close.
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
-	select {
-	case <-proxyResolverBuilt:
-	case <-ctx.Done():
-		t.Fatalf("Context timed out waiting for proxy resolver to be built.")
-	}
+	mustBuildResolver(ctx, t, proxyResolverBuilt)
 	// Closing the delegating resolver will block until the test writes to the
 	// unblockProxyResolverClose channel.
 	go dr.Close()
@@ -677,7 +707,7 @@ func (s) TestDelegatingResolverUpdateStateFromResolveNow(t *testing.T) {
 
 	target := targetResolver.Scheme() + ":///" + "ignored"
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
+	proxyResolver, proxyResolverBuilt := setupDNS(t)
 
 	tcc, _, _ := createTestResolverClientConn(t)
 	tcc.UpdateStateF = func(resolver.State) error {
@@ -692,14 +722,18 @@ func (s) TestDelegatingResolverUpdateStateFromResolveNow(t *testing.T) {
 		Endpoints: []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: "1.1.1.1"}}}},
 	})
 
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Wait for the proxy resolver to be built before calling UpdateState.
+	mustBuildResolver(ctx, t, proxyResolverBuilt)
+
 	// Updating the channel will result in an error being returned. The
 	// delegating resolver should call call "ResolveNow" on the target resolver.
 	proxyResolver.UpdateState(resolver.State{
 		Endpoints: []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: "1.1.1.1"}}}},
 	})
 
-	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-	defer cancel()
 	select {
 	case <-targetResolverCalled:
 	case <-ctx.Done():
@@ -737,12 +771,9 @@ func (s) TestDelegatingResolverResolveNow(t *testing.T) {
 
 	target := targetResolver.Scheme() + ":///" + "ignored"
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
+	proxyResolver, proxyResolverBuilt := setupDNS(t)
+
 	proxyResolverCalled := make(chan struct{})
-	proxyResolverBuilt := make(chan struct{})
-	proxyResolver.BuildCallback = func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) {
-		close(proxyResolverBuilt)
-	}
 	proxyResolver.ResolveNowCallback = func(resolver.ResolveNowOptions) {
 		// Updating the resolver state should not deadlock.
 		proxyResolver.CC().UpdateState(resolver.State{
@@ -770,12 +801,7 @@ func (s) TestDelegatingResolverResolveNow(t *testing.T) {
 		t.Fatalf("context timed out waiting for targetResolver.ResolveNow() to be called.")
 	}
 
-	// Wait for proxy resolver to be built.
-	select {
-	case <-proxyResolverBuilt:
-	case <-ctx.Done():
-		t.Fatalf("Timeout when waiting for proxy resolver to be built")
-	}
+	mustBuildResolver(ctx, t, proxyResolverBuilt)
 
 	dr.ResolveNow(resolver.ResolveNowOptions{})
 
@@ -821,11 +847,7 @@ func (s) TestDelegatingResolverForNonTCPTarget(t *testing.T) {
 	targetResolver := manual.NewBuilderWithScheme("test")
 	target := targetResolver.Scheme() + ":///" + targetTestAddr
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
-	proxyBuildCalled := make(chan struct{})
-	proxyResolver.BuildCallback = func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) {
-		close(proxyBuildCalled)
-	}
+	_, proxyResolverBuilt := setupDNS(t)
 
 	tcc, stateCh, _ := createTestResolverClientConn(t)
 	if _, err := delegatingresolver.New(resolver.Target{URL: *testutils.MustParseURL(target)}, tcc, resolver.BuildOptions{}, targetResolver, false); err != nil {
@@ -850,7 +872,7 @@ func (s) TestDelegatingResolverForNonTCPTarget(t *testing.T) {
 	// Verify that the delegating resolver doesn't call proxy resolver's
 	// UpdateState since we have no tcp address
 	select {
-	case <-proxyBuildCalled:
+	case <-proxyResolverBuilt:
 		t.Fatal("Unexpected call to proxy resolver update state")
 	case <-time.After(defaultTestShortTimeout):
 	}
@@ -900,7 +922,7 @@ func (s) TestDelegatingResolverForMixNetworkType(t *testing.T) {
 	targetResolver := manual.NewBuilderWithScheme("test")
 	target := targetResolver.Scheme() + ":///" + targetTestAddr
 	// Set up a manual DNS resolver to control the proxy address resolution.
-	proxyResolver := setupDNS(t)
+	proxyResolver, proxyResolverBuilt := setupDNS(t)
 
 	tcc, stateCh, _ := createTestResolverClientConn(t)
 	if _, err := delegatingresolver.New(resolver.Target{URL: *testutils.MustParseURL(target)}, tcc, resolver.BuildOptions{}, targetResolver, false); err != nil {
@@ -920,6 +942,11 @@ func (s) TestDelegatingResolverForMixNetworkType(t *testing.T) {
 	case <-time.After(defaultTestShortTimeout):
 	}
 
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Wait for the proxy resolver to be built before calling UpdateState.
+	mustBuildResolver(ctx, t, proxyResolverBuilt)
 	proxyResolver.UpdateState(resolver.State{
 		Addresses:     []resolver.Address{{Addr: envProxyAddr}},
 		ServiceConfig: &serviceconfig.ParseResult{},
@@ -928,8 +955,8 @@ func (s) TestDelegatingResolverForMixNetworkType(t *testing.T) {
 	var gotState resolver.State
 	select {
 	case gotState = <-stateCh:
-	case <-time.After(defaultTestTimeout):
-		t.Fatal("Timeout when waiting for a state update from the delegating resolver")
+	case <-ctx.Done():
+		t.Fatal("Context timed out when waiting for a state update from the delegating resolver")
 	}
 	wantState := resolver.State{
 		Addresses:     []resolver.Address{nonTCPAddr, proxyAddressWithTargetAttribute(envProxyAddr, resolvedTargetTestAddr2)},