Golang testutils.WithServer函数代码示例

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

	s1Opts := &testutils.ChannelOpts{ServiceName: "s1"}
	require.Nil(t, testutils.WithServer(s1Opts, func(ch1 *Channel, hostPort1 string) {
		s2Opts := &testutils.ChannelOpts{ServiceName: "s2"}
		require.Nil(t, testutils.WithServer(s2Opts, func(ch2 *Channel, hostPort2 string) {
			ch1.Register(raw.Wrap(newTestHandler(t)), "echo")
			ch2.Register(raw.Wrap(newTestHandler(t)), "echo")

			// We need the servers to have their peers set before making outgoing calls
			// for the outgoing calls to contain the correct peerInfo.
			require.True(t, testutils.WaitFor(time.Second, func() bool {
				return !ch1.PeerInfo().IsEphemeral() && !ch2.PeerInfo().IsEphemeral()
			}))

			outbound, err := ch1.BeginCall(ctx, hostPort2, "s2", "echo", nil)
			require.NoError(t, err)
			outboundConn, outboundNetConn := OutboundConnection(outbound)

			// Try to make another call at the same time, should reuse the same connection.
			outbound2, err := ch1.BeginCall(ctx, hostPort2, "s2", "echo", nil)
			require.NoError(t, err)
			outbound2Conn, _ := OutboundConnection(outbound)
			assert.Equal(t, outboundConn, outbound2Conn)

			// When ch2 tries to call ch1, it should reuse the inbound connection from ch1.
			outbound3, err := ch2.BeginCall(ctx, hostPort1, "s1", "echo", nil)
			require.NoError(t, err)
			_, outbound3NetConn := OutboundConnection(outbound3)
			assert.Equal(t, outboundNetConn.RemoteAddr(), outbound3NetConn.LocalAddr())
			assert.Equal(t, outboundNetConn.LocalAddr(), outbound3NetConn.RemoteAddr())

			// Ensure all calls can complete in parallel.
			var wg sync.WaitGroup
			for _, call := range []*OutboundCall{outbound, outbound2, outbound3} {
				wg.Add(1)
				go func(call *OutboundCall) {
					defer wg.Done()
					resp1, resp2, _, err := raw.WriteArgs(call, []byte("arg2"), []byte("arg3"))
					require.NoError(t, err)
					assert.Equal(t, resp1, []byte("arg2"), "result does match argument")
					assert.Equal(t, resp2, []byte("arg3"), "result does match argument")
				}(call)
			}
			wg.Wait()
		}))
	}))
}

func TestContextBuilder(t *testing.T) {
	ctx, cancel := tchannel.NewContextBuilder(time.Second).SetShardKey("shard").Build()
	defer cancel()

	var called bool
	testutils.WithServer(nil, func(ch *tchannel.Channel, hostPort string) {
		peerInfo := ch.PeerInfo()

		testutils.RegisterFunc(t, ch, "SecondService::Echo", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			call := tchannel.CurrentCall(ctx)
			assert.Equal(t, peerInfo.ServiceName, call.CallerName(), "unexpected caller name")
			assert.Equal(t, "shard", call.ShardKey(), "unexpected shard key")
			assert.Equal(t, tchannel.Thrift, args.Format)
			called = true
			return nil, errors.New("err")
		})

		client := NewClient(ch, ch.PeerInfo().ServiceName, &ClientOptions{
			HostPort: peerInfo.HostPort,
		})
		secondClient := gen.NewTChanSecondServiceClient(client)
		secondClient.Echo(ctx, "asd")
		assert.True(t, called, "test not called")
	})
}

func TestRoundTrip(t *testing.T) {
	require.Nil(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		handler := newTestHandler(t)
		ch.Register(raw.Wrap(handler), "echo")

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

		call, err := ch.BeginCall(ctx, hostPort, testServiceName, "echo", &CallOptions{Format: JSON})
		require.NoError(t, err)

		require.NoError(t, NewArgWriter(call.Arg2Writer()).Write(testArg2))
		require.NoError(t, NewArgWriter(call.Arg3Writer()).Write(testArg3))

		var respArg2 []byte
		require.NoError(t, NewArgReader(call.Response().Arg2Reader()).Read(&respArg2))
		assert.Equal(t, testArg2, []byte(respArg2))

		var respArg3 []byte
		require.NoError(t, NewArgReader(call.Response().Arg3Reader()).Read(&respArg3))
		assert.Equal(t, testArg3, []byte(respArg3))

		assert.Equal(t, JSON, handler.format)
		assert.Equal(t, testServiceName, handler.caller)
		assert.Equal(t, JSON, call.Response().Format(), "response Format should match request Format")
	}))
}

func TestTracingPropagates(t *testing.T) {
	require.Nil(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		handler := &traceHandler{t: t, ch: ch}
		json.Register(ch, map[string]interface{}{
			"call": handler.call,
		}, handler.onError)

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

		peer := ch.Peers().GetOrAdd(ch.PeerInfo().HostPort)
		var response TracingResponse
		require.NoError(t, json.CallPeer(ctx, peer, ch.PeerInfo().ServiceName, "call", &TracingRequest{
			ForwardCount: 1,
		}, &response))

		clientSpan := CurrentSpan(ctx)
		require.NotNil(t, clientSpan)
		assert.Equal(t, uint64(0), clientSpan.ParentID())

		assert.Equal(t, uint64(0), clientSpan.ParentID())
		assert.NotEqual(t, uint64(0), clientSpan.TraceID())
		assert.Equal(t, clientSpan.TraceID(), response.TraceID)
		assert.Equal(t, clientSpan.SpanID(), response.ParentID)
		assert.Equal(t, response.TraceID, response.SpanID, "traceID = spanID for root span")

		nestedResponse := response.Child
		require.NotNil(t, nestedResponse)
		assert.Equal(t, clientSpan.TraceID(), nestedResponse.TraceID)
		assert.Equal(t, response.SpanID, nestedResponse.ParentID)
		assert.NotEqual(t, response.SpanID, nestedResponse.SpanID)
	}))
}

func TestNoTimeout(t *testing.T) {
	require.Nil(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		ch.Register(raw.Wrap(newTestHandler(t)), "Echo")

		ctx := context.Background()
		_, _, _, err := raw.Call(ctx, ch, hostPort, "svc", "Echo", []byte("Headers"), []byte("Body"))
		require.NotNil(t, err)
		assert.Equal(t, ErrTimeoutRequired, err)
	}))
}

func TestBadRequest(t *testing.T) {
	require.Nil(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		ctx, cancel := NewContext(time.Second * 5)
		defer cancel()

		_, _, _, err := raw.Call(ctx, ch, hostPort, "Nowhere", "Noone", []byte("Headers"), []byte("Body"))
		require.NotNil(t, err)
		assert.Equal(t, ErrCodeBadRequest, GetSystemErrorCode(err))
	}))
}

func TestPing(t *testing.T) {
	require.Nil(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		ctx, cancel := NewContext(time.Second * 5)
		defer cancel()

		clientCh, err := testutils.NewClient(nil)
		require.NoError(t, err)
		require.NoError(t, clientCh.Ping(ctx, hostPort))
	}))
}

func TestServerBusy(t *testing.T) {
	require.Nil(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		ch.Register(raw.Wrap(newTestHandler(t)), "busy")

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

		_, _, _, err := raw.Call(ctx, ch, hostPort, testServiceName, "busy", []byte("Arg2"), []byte("Arg3"))
		require.NotNil(t, err)
		assert.Equal(t, ErrCodeBusy, GetSystemErrorCode(err), "err: %v", err)
	}))
}

func TestTimeout(t *testing.T) {
	require.Nil(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		ch.Register(raw.Wrap(newTestHandler(t)), "timeout")

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

		_, _, _, err := raw.Call(ctx, ch, hostPort, testServiceName, "timeout", []byte("Arg2"), []byte("Arg3"))

		// TODO(mmihic): Maybe translate this into ErrTimeout (or vice versa)?
		assert.Equal(t, context.DeadlineExceeded, err)
	}))
}

func TestDefaultFormat(t *testing.T) {
	require.Nil(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		handler := newTestHandler(t)
		ch.Register(raw.Wrap(handler), "echo")

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

		arg2, arg3, resp, err := raw.Call(ctx, ch, hostPort, testServiceName, "echo", testArg2, testArg3)
		require.Nil(t, err)

		require.Equal(t, testArg2, arg2)
		require.Equal(t, testArg3, arg3)
		require.Equal(t, Raw, handler.format)
		assert.Equal(t, Raw, resp.Format(), "response Format should match request Format")
	}))
}

func TestStatsCalls(t *testing.T) {
	serverStats := newRecordingStatsReporter()
	serverOpts := &testutils.ChannelOpts{
		StatsReporter: serverStats,
	}
	require.NoError(t, testutils.WithServer(serverOpts, func(serverCh *Channel, hostPort string) {
		handler := raw.Wrap(newTestHandler(t))
		serverCh.Register(handler, "echo")
		serverCh.Register(handler, "app-error")

		clientStats := newRecordingStatsReporter()
		ch, err := testutils.NewClient(&testutils.ChannelOpts{StatsReporter: clientStats})
		require.NoError(t, err)

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

		_, _, _, err = raw.Call(ctx, ch, hostPort, testServiceName, "echo", []byte("Headers"), []byte("Body"))
		require.NoError(t, err)

		_, _, resp, err := raw.Call(ctx, ch, hostPort, testServiceName, "app-error", nil, nil)
		require.NoError(t, err)
		require.True(t, resp.ApplicationError(), "expected application error")

		outboundTags := tagsForOutboundCall(serverCh, ch, "echo")
		clientStats.Expected.IncCounter("outbound.calls.send", outboundTags, 1)
		clientStats.Expected.IncCounter("outbound.calls.success", outboundTags, 1)
		outboundTags["target-endpoint"] = "app-error"
		clientStats.Expected.IncCounter("outbound.calls.send", outboundTags, 1)
		clientStats.Expected.IncCounter("outbound.calls.app-errors", outboundTags, 1)

		inboundTags := tagsForInboundCall(serverCh, ch, "echo")
		serverStats.Expected.IncCounter("inbound.calls.recvd", inboundTags, 1)
		serverStats.Expected.IncCounter("inbound.calls.success", inboundTags, 1)
		inboundTags["endpoint"] = "app-error"
		serverStats.Expected.IncCounter("inbound.calls.recvd", inboundTags, 1)
		serverStats.Expected.IncCounter("inbound.calls.app-errors", inboundTags, 1)

		clientStats.ValidateCounters(t)
		serverStats.ValidateCounters(t)
	}))
}

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

	// Note: This test leaks a message exchange due to the modification of IDs in the relay.
	require.NoError(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		gotCall := make(chan struct{})
		unblock := make(chan struct{})

		testutils.RegisterFunc(t, ch, "blocked", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			gotCall <- struct{}{}
			<-unblock
			return &raw.Res{}, nil
		})

		relayFunc := func(outgoing bool, frame *Frame) *Frame {
			if outgoing && frame.Header.ID == 2 {
				frame.Header.ID = 3
			}
			return frame
		}

		relayHostPort, closeRelay := testutils.FrameRelay(t, hostPort, relayFunc)
		defer closeRelay()

		go func() {
			// This call will block until we close unblock.
			raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
		}()

		// Wait for the first call to be received by the server
		<-gotCall

		// Make a new call, which should fail
		_, _, _, err := raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
		assert.Error(t, err, "Expect error")
		assert.True(t, strings.Contains(err.Error(), "already active"),
			"expected already active error, got %v", err)

		close(unblock)
	}))
}

func TestLargeRequest(t *testing.T) {
	if !*flagStressTest {
		t.Skip("skipping as stress tests are not enabled")
	}

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

		maxRequestSize = 1 * GB
	)

	require.NoError(t, testutils.WithServer(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 := makeData(reqSize)
			arg2 := arg3[len(arg3)/2:]

			clientCh, err := testutils.NewClient(nil)
			require.NoError(t, err, "new client failed")

			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 TestFragmentation(t *testing.T) {
	require.Nil(t, testutils.WithServer(nil, func(ch *Channel, hostPort string) {
		ch.Register(raw.Wrap(newTestHandler(t)), "echo")

		arg2 := make([]byte, MaxFramePayloadSize*2)
		for i := 0; i < len(arg2); i++ {
			arg2[i] = byte('a' + (i % 10))
		}

		arg3 := make([]byte, MaxFramePayloadSize*3)
		for i := 0; i < len(arg3); i++ {
			arg3[i] = byte('A' + (i % 10))
		}

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

		respArg2, respArg3, _, err := raw.Call(ctx, ch, hostPort, testServiceName, "echo", arg2, arg3)
		require.NoError(t, err)
		assert.Equal(t, arg2, respArg2)
		assert.Equal(t, arg3, respArg3)
	}))
}

func TestStatsCalls(t *testing.T) {
	statsReporter := newRecordingStatsReporter()
	testOpts := &testutils.ChannelOpts{
		StatsReporter: statsReporter,
	}
	require.NoError(t, testutils.WithServer(testOpts, func(ch *Channel, hostPort string) {
		ch.Register(raw.Wrap(newTestHandler(t)), "echo")

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

		_, _, _, err := raw.Call(ctx, ch, hostPort, testServiceName, "echo", []byte("Headers"), []byte("Body"))
		require.NoError(t, err)

		_, _, _, err = raw.Call(ctx, ch, hostPort, testServiceName, "error", nil, nil)
		require.Error(t, err)

		host, err := os.Hostname()
		require.Nil(t, err)

		expectedTags := map[string]string{
			"app":             ch.PeerInfo().ProcessName,
			"host":            host,
			"service":         ch.PeerInfo().ServiceName,
			"target-service":  ch.PeerInfo().ServiceName,
			"target-endpoint": "echo",
		}
		statsReporter.Expected.IncCounter("outbound.calls.send", expectedTags, 1)
		statsReporter.Expected.IncCounter("outbound.calls.successful", expectedTags, 1)
		expectedTags["target-endpoint"] = "error"
		statsReporter.Expected.IncCounter("outbound.calls.send", expectedTags, 1)
		// TODO(prashant): Make the following stat work too.
		// statsReporter.Expected.IncCounter("outbound.calls.app-errors", expectedTags, 1)
		statsReporter.ValidateCounters(t)
	}))
}