Repository: prometheus/prometheus
File: storage/remote/dial_context_test.go (Go)

Code before:
func (lh *mockedLookupHost) LookupHost(context.Context, string) ([]string, error) {
	if lh.withErr {
		return nil, errMockLookupHost
	}
	return lh.result, nil
}

func TestDialContextWithRandomConnections(t *testing.T) {
	numberOfRuns := 2 * len(testLookupResult)
	var mdc *mockDialContext
	testCases := map[string]struct {
		addr                              string
		setup                             func() dialContextWithRoundRobinDNS
		check                             func()
		withLookupHostErr                 bool
		expectedDefaultContextArguments   []string
		unexpectedDefaultContextArguments []string
	}{
		"if address contains no port call default DealContext": {
			addr: testAddrWithoutPort,
			setup: func() dialContextWithRoundRobinDNS {
				mdc = newMockDialContext([]string{testAddrWithoutPort})
				return dialContextWithRoundRobinDNS{
					dialContext: mdc.dialContext,
					resolver:    &mockedLookupHost{withErr: false},
				}
			},
			check: func() {
				require.Equal(t, numberOfRuns, mdc.getCount(testAddrWithoutPort))
			},
		},
		"if lookup host returns error call default DealContext": {
			addr: testAddrWithPort,
			setup: func() dialContextWithRoundRobinDNS {
				mdc = newMockDialContext([]string{testAddrWithPort})
				return dialContextWithRoundRobinDNS{
					dialContext: mdc.dialContext,
					resolver:    &mockedLookupHost{withErr: true},
				}
			},
			check: func() {
				require.Equal(t, numberOfRuns, mdc.getCount(testAddrWithPort))
			},
		},
		"if lookup host is successful, shuffle results": {
			addr: testAddrWithPort,
			setup: func() dialContextWithRoundRobinDNS {
				mdc = newMockDialContext(testLookupResultWithPort)
				return dialContextWithRoundRobinDNS{
					dialContext: mdc.dialContext,
					resolver:    &mockedLookupHost{result: testLookupResult},

Review comment (refactor):
Worth extracting this part to a function?  It is repeated 3 times with slight differences, so pulling those into a single function-call might make it easier to follow.

Code after:
	withErr bool
	result  []string
}

func (lh *mockedLookupHost) LookupHost(context.Context, string) ([]string, error) {
	if lh.withErr {
		return nil, errMockLookupHost
	}
	return lh.result, nil
}

func createDialContextWithRoundRobinDNS(dialContext config.DialContextFunc, resolver hostResolver, r *rand.Rand) dialContextWithRoundRobinDNS {
	return dialContextWithRoundRobinDNS{
		dialContext: dialContext,
		resolver:    resolver,
		rand:        r,
	}
}

func TestDialContextWithRandomConnections(t *testing.T) {
	numberOfRuns := 2 * len(testLookupResult)
	var mdc *mockDialContext
	testCases := map[string]struct {
		addr  string
		setup func() dialContextWithRoundRobinDNS
		check func()
	}{
		"if address contains no port call default DealContext": {
			addr: testAddrWithoutPort,
			setup: func() dialContextWithRoundRobinDNS {
				mdc = newMockDialContext([]string{testAddrWithoutPort})
				return createDialContextWithRoundRobinDNS(mdc.dialContext, &mockedLookupHost{withErr: false}, rand.New(rand.NewSource(time.Now().Unix())))
			},
			check: func() {
				require.Equal(t, numberOfRuns, mdc.getCount(testAddrWithoutPort))
			},
		},
		"if lookup host returns error call default DealContext": {
			addr: testAddrWithPort,
			setup: func() dialContextWithRoundRobinDNS {
				mdc = newMockDialContext([]string{testAddrWithPort})
				return createDialContextWithRoundRobinDNS(mdc.dialContext, &mockedLookupHost{withErr: true}, rand.New(rand.NewSource(time.Now().Unix())))
			},
			check: func() {
				require.Equal(t, numberOfRuns, mdc.getCount(testAddrWithPort))
			},
		},
		"if lookup returns no addresses call default DealContext": {
			addr: testAddrWithPort,
			setup: func() dialContextWithRoundRobinDNS {
				mdc = newMockDialContext([]string{testAddrWithPort})