Golang testutils.NewClient函数代码示例

func BenchmarkBothSerial(b *testing.B) {
	serverAddr, err := setupBenchServer()
	require.NoError(b, err, "setupBenchServer failed")

	opts := testutils.NewOpts().SetFramePool(tchannel.NewSyncFramePool())
	clientCh := testutils.NewClient(b, opts)
	for _, addr := range serverAddr {
		clientCh.Peers().Add(addr)
	}

	thriftClient := thrift.NewClient(clientCh, "bench-server", nil)
	client := gen.NewTChanSecondServiceClient(thriftClient)
	ctx, cancel := thrift.NewContext(10 * time.Millisecond)
	client.Echo(ctx, "make connection")
	cancel()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ctx, cancel := thrift.NewContext(10 * time.Millisecond)
		defer cancel()

		_, err := client.Echo(ctx, "hello world")
		if err != nil {
			b.Errorf("Echo failed: %v", err)
		}
	}
}

func TestRequestSubChannel(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	tchan := testutils.NewServer(t, testutils.NewOpts().SetServiceName("svc1"))
	defer tchan.Close()

	clientCh := testutils.NewClient(t, nil)
	defer clientCh.Close()
	clientCh.Peers().Add(tchan.PeerInfo().HostPort)

	tests := []tchannel.Registrar{tchan, tchan.GetSubChannel("svc2"), tchan.GetSubChannel("svc3")}
	for _, ch := range tests {
		mockHandler := new(mocks.TChanSecondService)
		server := NewServer(ch)
		server.Register(gen.NewTChanSecondServiceServer(mockHandler))

		client := NewClient(clientCh, ch.ServiceName(), nil)
		secondClient := gen.NewTChanSecondServiceClient(client)

		echoArg := ch.ServiceName()
		echoRes := echoArg + "-echo"
		mockHandler.On("Echo", ctxArg(), echoArg).Return(echoRes, nil)
		res, err := secondClient.Echo(ctx, echoArg)
		assert.NoError(t, err, "Echo failed")
		assert.Equal(t, echoRes, res)
	}
}

func TestRequestStateRetry(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	server := testutils.NewServer(t, nil)
	defer server.Close()
	server.Register(raw.Wrap(newTestHandler(t)), "echo")

	client := testutils.NewClient(t, nil)
	defer client.Close()

	counter := 0
	sc := client.GetSubChannel(server.PeerInfo().ServiceName)
	err := client.RunWithRetry(ctx, func(ctx context.Context, rs *RequestState) error {
		defer func() { counter++ }()

		assert.Equal(t, counter, len(rs.SelectedPeers), "SelectedPeers should not be reused")

		if counter < 4 {
			client.Peers().Add(testutils.GetClosedHostPort(t))
		} else {
			client.Peers().Add(server.PeerInfo().HostPort)
		}

		_, err := raw.CallV2(ctx, sc, raw.CArgs{
			Operation:   "echo",
			CallOptions: &CallOptions{RequestState: rs},
		})
		return err
	})
	assert.NoError(t, err, "RunWithRetry should succeed")
	assert.Equal(t, 5, counter, "RunWithRetry should retry 5 times")
}

func TestGetPeerNoPeer(t *testing.T) {
	ch := testutils.NewClient(t, nil)
	defer ch.Close()
	peer, err := ch.Peers().Get(nil)
	assert.Equal(t, ErrNoPeers, err, "Empty peer list should return error")
	assert.Nil(t, peer, "should not return peer")
}

func TestRetryNetConnect(t *testing.T) {
	e := getTestErrors()
	ch := testutils.NewClient(t, nil)
	defer ch.Close()

	ctx, cancel := NewContext(time.Second)
	defer cancel()

	closedAddr := testutils.GetClosedHostPort(t)
	listenC, err := net.Listen("tcp", ":0")
	require.NoError(t, err, "Listen failed")
	defer listenC.Close()

	counter := 0
	f := func(ctx context.Context, rs *RequestState) error {
		counter++
		if !rs.HasRetries(e.Connection) {
			c, err := net.Dial("tcp", listenC.Addr().String())
			if err == nil {
				c.Close()
			}
			return err
		}

		_, err := net.Dial("tcp", closedAddr)
		return err
	}

	assert.NoError(t, ch.RunWithRetry(ctx, f), "RunWithRetry should succeed")
	assert.Equal(t, 5, counter, "RunWithRetry should have run f 5 times")
}

func TestLargeRequest(t *testing.T) {
	CheckStress(t)

	const (
		KB = 1024
		MB = 1024 * KB
		GB = 1024 * MB

		maxRequestSize = 1 * GB
	)

	WithVerifiedServer(t, nil, func(serverCh *Channel, hostPort string) {
		serverCh.Register(raw.Wrap(newTestHandler(t)), "echo")

		for reqSize := 2; reqSize <= maxRequestSize; reqSize *= 2 {
			log.Printf("reqSize = %v", reqSize)
			arg3 := testutils.RandBytes(reqSize)
			arg2 := testutils.RandBytes(reqSize / 2)

			clientCh := testutils.NewClient(t, nil)
			ctx, cancel := NewContext(time.Second * 30)
			rArg2, rArg3, _, err := raw.Call(ctx, clientCh, hostPort, serverCh.PeerInfo().ServiceName, "echo", arg2, arg3)
			require.NoError(t, err, "Call failed")

			if !bytes.Equal(arg2, rArg2) {
				t.Errorf("echo arg2 mismatch")
			}
			if !bytes.Equal(arg3, rArg3) {
				t.Errorf("echo arg3 mismatch")
			}
			cancel()
		}
	})
}

func BenchmarkCallsConcurrent(b *testing.B) {
	const numWorkers = 5

	serverCh, svcName, svcHostPort := setupServer(b)
	defer serverCh.Close()

	var wg sync.WaitGroup
	inCh := make(chan struct{})
	for i := 0; i < numWorkers; i++ {
		go func() {
			clientCh := testutils.NewClient(b, nil)
			defer clientCh.Close()

			for range inCh {
				ctx, cancel := NewContext(time.Second)

				_, _, _, err := raw.Call(ctx, clientCh, svcHostPort, svcName, "echo", []byte("data111"), []byte("data222"))
				assert.NoError(b, err)

				cancel()
				wg.Done()
			}
		}()
	}

	for i := 0; i < b.N; i++ {
		wg.Add(1)
		inCh <- struct{}{}
	}

	wg.Wait()
	close(inCh)
}

func TestNetDialTimeout(t *testing.T) {
	// timeoutHostPort uses a blackholed address (RFC 6890) with a port
	// reserved for documentation. This address should always cause a timeout.
	const timeoutHostPort = "192.18.0.254:44444"
	timeoutPeriod := testutils.Timeout(50 * time.Millisecond)

	client := testutils.NewClient(t, nil)
	defer client.Close()

	started := time.Now()
	ctx, cancel := NewContext(timeoutPeriod)
	defer cancel()

	err := client.Ping(ctx, timeoutHostPort)
	if !assert.Error(t, err, "Ping to blackhole address should fail") {
		return
	}

	if strings.Contains(err.Error(), "network is unreachable") {
		t.Skipf("Skipping test, as network interface may not be available")
	}

	d := time.Since(started)
	assert.Equal(t, ErrTimeout, err, "Ping expected to fail with timeout")
	assert.True(t, d >= timeoutPeriod, "Timeout should take more than %v, took %v", timeoutPeriod, d)
}

func TestInboundEphemeralPeerRemoved(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	WithVerifiedServer(t, nil, func(ch *Channel, hostPort string) {
		client := testutils.NewClient(t, nil)
		assert.NoError(t, client.Ping(ctx, hostPort), "Ping to server failed")

		// Server should have a host:port in the root peers for the client.
		var clientHP string
		peers := ch.RootPeers().Copy()
		for k := range peers {
			clientHP = k
		}

		// Close the connection, which should remove the peer from the server channel.
		client.Close()
		runtime.Gosched()
		assert.Equal(t, ChannelClosed, client.State(), "Client should be closed")

		// Wait for the channel to see the connection as closed and update the peer list.
		time.Sleep(time.Millisecond)

		_, ok := ch.RootPeers().Get(clientHP)
		assert.False(t, ok, "server's root peers should remove peer for client on close")
	})
}

func TestInboundEphemeralPeerRemoved(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	// No relay, since we look for the exact host:port in peer lists.
	opts := testutils.NewOpts().NoRelay()
	WithVerifiedServer(t, opts, func(ch *Channel, hostPort string) {
		client := testutils.NewClient(t, nil)
		assert.NoError(t, client.Ping(ctx, hostPort), "Ping to server failed")

		// Server should have a host:port in the root peers for the client.
		var clientHP string
		peers := ch.RootPeers().Copy()
		for k := range peers {
			clientHP = k
		}

		waitTillInboundEmpty(t, ch, clientHP, func() {
			client.Close()
		})
		assert.Equal(t, ChannelClosed, client.State(), "Client should be closed")

		_, ok := ch.RootPeers().Get(clientHP)
		assert.False(t, ok, "server's root peers should remove peer for client on close")
	})
}

func TestPeerSelectionRanking(t *testing.T) {
	const numPeers = 10
	const numIterations = 1000

	// Selected is a map from rank -> [peer, count]
	// It tracks how often a peer gets selected at a specific rank.
	selected := make([]map[string]int, numPeers)
	for i := 0; i < numPeers; i++ {
		selected[i] = make(map[string]int)
	}

	for i := 0; i < numIterations; i++ {
		ch := testutils.NewClient(t, nil)
		defer ch.Close()
		ch.SetRandomSeed(int64(i * 100))

		for i := 0; i < numPeers; i++ {
			hp := fmt.Sprintf("127.0.0.1:60%v", i)
			ch.Peers().Add(hp)
		}

		for i := 0; i < numPeers; i++ {
			peer, err := ch.Peers().Get(nil)
			require.NoError(t, err, "Peers.Get failed")
			selected[i][peer.HostPort()]++
		}
	}

	for _, m := range selected {
		testDistribution(t, m, 50, 150)
	}
}

func TestCloseAfterTimeout(t *testing.T) {
	WithVerifiedServer(t, nil, func(ch *Channel, hostPort string) {
		testHandler := onErrorTestHandler{newTestHandler(t), func(_ context.Context, err error) {}}
		ch.Register(raw.Wrap(testHandler), "block")

		ctx, cancel := NewContext(10 * time.Millisecond)
		defer cancel()

		// Make a call, wait for it to timeout.
		clientCh, err := testutils.NewClient(nil)
		require.NoError(t, err, "NewClient failed")
		peerInfo := ch.PeerInfo()
		_, _, _, err = raw.Call(ctx, clientCh, peerInfo.HostPort, peerInfo.ServiceName, "block", nil, nil)
		require.Error(t, err, "Expected call to timeout")

		// The client channel should also close immediately.
		clientCh.Close()
		runtime.Gosched()
		assert.Equal(t, ChannelClosed, clientCh.State())
		assert.True(t, clientCh.Closed(), "Channel should be closed")

		// Unblock the testHandler so that a goroutine isn't leaked.
		<-testHandler.blockErr
	})
	VerifyNoBlockedGoroutines(t)
}

func setupProxy(t *testing.T, tchanAddr string) (string, func()) {
	mux := http.NewServeMux()
	mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// You get /proxy/host:port/rest/of/the/path
		parts := strings.SplitN(r.URL.Path, "/", 4)
		r.URL.Host = parts[2]
		r.URL.Scheme = "http"
		r.URL.Path = parts[3]

		ch := testutils.NewClient(t, nil)
		ctx, cancel := tchannel.NewContext(time.Second)
		defer cancel()

		call, err := ch.BeginCall(ctx, tchanAddr, "test", "http", nil)
		require.NoError(t, err, "BeginCall failed")

		require.NoError(t, WriteRequest(call, r), "WriteRequest failed")
		resp, err := ReadResponse(call.Response())
		require.NoError(t, err, "Read response failed")

		for k, vs := range resp.Header {
			for _, v := range vs {
				w.Header().Add(k, v)
			}
		}
		w.WriteHeader(resp.StatusCode)

		_, err = io.Copy(w, resp.Body)
		assert.NoError(t, err, "io.Copy failed")
		err = resp.Body.Close()
		assert.NoError(t, err, "Close Response Body failed")
	}))
	return setupHTTP(t, mux)
}

func TestRetrySubContextTimeoutPerAttempt(t *testing.T) {
	e := getTestErrors()
	ctx, cancel := NewContextBuilder(time.Second).
		SetTimeoutPerAttempt(time.Millisecond).Build()
	defer cancel()

	ch := testutils.NewClient(t, nil)
	defer ch.Close()

	var lastDeadline time.Time

	counter := 0
	ch.RunWithRetry(ctx, func(sctx context.Context, _ *RequestState) error {
		counter++

		assert.NotEqual(t, ctx, sctx, "Sub-context should be different")
		deadline, _ := sctx.Deadline()
		assert.True(t, deadline.After(lastDeadline), "Deadline is invalid")
		lastDeadline = deadline

		overallDeadline, _ := ctx.Deadline()
		assert.True(t, overallDeadline.After(deadline), "Deadline is invalid")

		return e.Busy
	})
	assert.Equal(t, 5, counter, "RunWithRetry did not run f enough times")
}

func TestStopMultiple(t *testing.T) {
	for _, getOutbound := range newOutbounds {
		out := getOutbound(testutils.NewClient(t, &testutils.ChannelOpts{
			ServiceName: "caller",
		}), "localhost:4040")
		// TODO: If we change Start() to establish a connection to the host, this
		// hostport will have to be changed to a real server.

		err := out.Start(transport.NoDeps)
		require.NoError(t, err)

		var wg sync.WaitGroup
		signal := make(chan struct{})

		for i := 0; i < 10; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				<-signal

				err := out.Stop()
				assert.NoError(t, err)
			}()
		}
		close(signal)
		wg.Wait()
	}
}