Using Toxiproxy for resilliency tests

grpc_test.go

import (
	"net"
	"strconv"
	"testing"
	"time"

	"github.com/Shopify/toxiproxy"
	toxiclient "github.com/Shopify/toxiproxy/client"
)

func TestMyGRPCService(t *testing.T) {
	// start the gRPC service
	svc := NewMyGRPCService()
	addr, serverDestroyFunc := startMyGRPCServiceServer(t, svc)
	defer serverDestroyFunc()

	// start toxiproxy instance that proxies traffic to the gRPC service
	proxyAddr, proxy, proxyClient := startToxiProxy(t, addr)
	defer proxy.Delete()

	// start the gRPC client that communicates with the service via toxiproxy
	svcClient, svcClientDestroyFunc := createMyGRPCServiceClient(t, proxyAddr)
	defer svcClientDestroyFunc()

	testCases := []struct {
		name    string
		toxicFn func() (func(), error)
	}{
		{
			name: "baseline",
			toxicFn: func() (func(), error) {
				return func() {}, nil
			},
		},
		{
			name: "wifi_connection",
			toxicFn: func() (func(), error) {
				t1, err := proxy.AddToxic("", "latency", "", 1.0, toxiclient.Attributes(map[string]interface{}{"latency": 40, "jitter": 10}))
				if err != nil {
					return nil, err
				}

				t2, err := proxy.AddToxic("", "bandwidth", "", 1.0, toxiclient.Attributes(map[string]interface{}{"rate": 30000}))
				if err != nil {
					return nil, err
				}

				return func() {
					proxy.RemoveToxic(t1.Name)
					proxy.RemoveToxic(t2.Name)
				}, nil
			},
		},
		{
			name: "3g_connection",
			toxicFn: func() (func(), error) {
				t1, err := proxy.AddToxic("", "latency", "", 1.0, toxiclient.Attributes(map[string]interface{}{"latency": 250, "jitter": 50}))
				if err != nil {
					return nil, err
				}

				t2, err := proxy.AddToxic("", "bandwidth", "", 1.0, toxiclient.Attributes(map[string]interface{}{"rate": 750}))
				if err != nil {
					return nil, err
				}

				return func() {
					proxy.RemoveToxic(t1.Name)
					proxy.RemoveToxic(t2.Name)
				}, nil
			},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			if testing.Short() {
				t.Skip("Skipping toxiproxy tests")
			}

			if err := proxyClient.ResetState(); err != nil {
				failNow(t, "Failed to reset toxiproxy state: %+v", err)
			}

			cleanupFn, err := tc.toxicFn()
			if err != nil {
				failNow(t, "Failed to add toxic: %+v", err)
			}
			defer cleanupFn()

			t.Run("method_x", testMethodX(svcClient))
			t.Run("method_y", testMethodY(svcClient))
			// other gRPC method tests
		})
	}
}

func testMethodX(client *MyGRPCServiceClient) func(*testing.T) {
	return func(t *testing.T) {
		// test method X
	}
}

func testMethodY(client *MyGRPCServiceClient) func(*testing.T) {
	return func(t *testing.T) {
		// test method Y
	}
}

func startToxiProxy(t *testing.T, svcAddr string) (string, *toxiclient.Proxy, *toxiclient.Client) {
	t.Helper()

	toxiAPIPort := getFreePort(t)
	proxyPort := getFreePort(t)

	// launch toxiproxy instance
	toxiServer := toxiproxy.NewServer()
	go toxiServer.Listen("localhost", toxiAPIPort)

	// create toxiproxy client
	proxyClient := toxiclient.NewClient(net.JoinHostPort("localhost", toxiAPIPort))
	proxyAddr := net.JoinHostPort("localhost", proxyPort)

	// wait for the toxiproxy instance to become ready
	for i := 0; i < 5; i++ {
		proxy, err := proxyClient.CreateProxy("myGRPCService", proxyAddr, svcAddr)
		if err != nil {
			if i < 4 {
				time.Sleep(50 * time.Millisecond)
				continue
			}
			t.Fatalf("Failed to start proxy: %+v", err)
		}

		return proxyAddr, proxy, proxyClient
	}

	return "", nil, nil
}

func getFreePort(t *testing.T) string {
	t.Helper()

	lis, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatal(err)
	}
	defer lis.Close()

	return strconv.Itoa(lis.Addr().(*net.TCPAddr).Port)
}

func failNow(t *testing.T, msg string, args ...interface{}) {
	t.Helper()
	t.Errorf(msg, args...)
	t.FailNow()
}