From e5b33f5d6bed6b3788e919c8d0a1289f49f393e1 Mon Sep 17 00:00:00 2001 From: Dean Sheather Date: Fri, 23 Feb 2024 10:40:15 +0000 Subject: [PATCH 1/2] chore: fix wsproxy test flake --- enterprise/wsproxy/wsproxy_test.go | 34 ++++++++++++++++++++---------- 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/enterprise/wsproxy/wsproxy_test.go b/enterprise/wsproxy/wsproxy_test.go index 0d440165dfb16..65b6f706a79f7 100644 --- a/enterprise/wsproxy/wsproxy_test.go +++ b/enterprise/wsproxy/wsproxy_test.go @@ -344,7 +344,7 @@ func TestDERPEndToEnd(t *testing.T) { Name: "best-proxy", }) - // Wait for the proxy to become healthy. + // Wait for the proxy to become healthy and appear in the DERP map. require.Eventually(t, func() bool { healthCtx := testutil.Context(t, testutil.WaitLong) err := api.ProxyHealth.ForceUpdate(healthCtx) @@ -367,17 +367,29 @@ func TestDERPEndToEnd(t *testing.T) { return true }, testutil.WaitLong, testutil.IntervalMedium) - // Swap out the DERPMapper for a fake one that only returns the proxy. This - // allows us to force the agent to pick the proxy as its preferred region. - oldDERPMapper := *api.AGPL.DERPMapper.Load() - newDERPMapper := func(derpMap *tailcfg.DERPMap) *tailcfg.DERPMap { - derpMap = oldDERPMapper(derpMap) - // Strip everything but the proxy, which is region ID 10001. - derpMap.Regions = map[int]*tailcfg.DERPRegion{ - 10001: derpMap.Regions[10001], + // Wait until the proxy appears in the DERP map, and then swap out the DERP + // map for one that only contains the proxy region. This allows us to force + // the agent to pick the proxy as its preferred region. + var proxyOnlyDERPMap *tailcfg.DERPMap + require.Eventually(t, func() bool { + derpMap := api.AGPL.DERPMap() + if derpMap == nil { + return false + } + if _, ok := derpMap.Regions[10001]; !ok { + return false } - derpMap.OmitDefaultRegions = true - return derpMap + + // Make a DERP map that only contains the proxy region. + proxyOnlyDERPMap = derpMap.Clone() + proxyOnlyDERPMap.Regions = map[int]*tailcfg.DERPRegion{ + 10001: proxyOnlyDERPMap.Regions[10001], + } + proxyOnlyDERPMap.OmitDefaultRegions = true + return true + }, testutil.WaitLong, testutil.IntervalMedium) + newDERPMapper := func(derpMap *tailcfg.DERPMap) *tailcfg.DERPMap { + return proxyOnlyDERPMap } api.AGPL.DERPMapper.Store(&newDERPMapper) From da5ad69aff8ccd303e3fc7265e0836ef0d093bb2 Mon Sep 17 00:00:00 2001 From: Dean Sheather Date: Fri, 23 Feb 2024 10:40:51 +0000 Subject: [PATCH 2/2] fixup! chore: fix wsproxy test flake --- enterprise/wsproxy/wsproxy_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/enterprise/wsproxy/wsproxy_test.go b/enterprise/wsproxy/wsproxy_test.go index 65b6f706a79f7..83f9321bdeaa8 100644 --- a/enterprise/wsproxy/wsproxy_test.go +++ b/enterprise/wsproxy/wsproxy_test.go @@ -344,7 +344,7 @@ func TestDERPEndToEnd(t *testing.T) { Name: "best-proxy", }) - // Wait for the proxy to become healthy and appear in the DERP map. + // Wait for the proxy to become healthy. require.Eventually(t, func() bool { healthCtx := testutil.Context(t, testutil.WaitLong) err := api.ProxyHealth.ForceUpdate(healthCtx)