Add grpc-gateway tests for all APIv3 methods

cmd/query/app/apiv3/grpc_gateway_test.go

@@ -15,18 +15,25 @@
 package apiv3
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"net"
 	"net/http"
+	"net/url"
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
+	gogojsonpb "github.com/gogo/protobuf/jsonpb"
+	gogoproto "github.com/gogo/protobuf/proto"
 	"github.com/gorilla/mux"
-	"github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
@@ -39,23 +46,36 @@ import (
 	"github.com/jaegertracing/jaeger/pkg/tenancy"
 	"github.com/jaegertracing/jaeger/proto-gen/api_v3"
 	dependencyStoreMocks "github.com/jaegertracing/jaeger/storage/dependencystore/mocks"
+	"github.com/jaegertracing/jaeger/storage/spanstore"
 	spanstoremocks "github.com/jaegertracing/jaeger/storage/spanstore/mocks"
 )
 
-var testCertKeyLocation = "../../../../pkg/config/tlscfg/testdata/"
+const (
+	testCertKeyLocation = "../../../../pkg/config/tlscfg/testdata/"
+	snapshotLocation    = "./snapshots/"
+)
+
+// Snapshots can be regenerated via:
+//
+//	REGENERATE_SNAPSHOTS=true go test -v ./cmd/query/app/apiv3/...
+var regenerateSnapshots = os.Getenv("REGENERATE_SNAPSHOTS") == "true"
 
 type testGateway struct {
 	reader *spanstoremocks.Reader
 	url    string
 }
 
+type gatewayRequest struct {
+	url          string
+	setupRequest func(*http.Request)
+}
+
 func setupGRPCGateway(
 	t *testing.T,
 	basePath string,
 	serverTLS, clientTLS *tlscfg.Options,
 	tenancyOptions tenancy.Options,
 ) *testGateway {
-	// *spanstoremocks.Reader, net.Listener, *grpc.Server, context.CancelFunc, *http.Server
 	gw := &testGateway{
 		reader: &spanstoremocks.Reader{},
 	}
@@ -123,6 +143,64 @@ func setupGRPCGateway(
 	return gw
 }
 
+func (gw *testGateway) execRequest(t *testing.T, gwReq *gatewayRequest) ([]byte, int) {
+	req, err := http.NewRequest(http.MethodGet, gw.url+gwReq.url, nil)
+	require.NoError(t, err)
+	req.Header.Set("Content-Type", "application/json")
+	gwReq.setupRequest(req)
+	response, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	body, err := io.ReadAll(response.Body)
+	require.NoError(t, err)
+	require.NoError(t, response.Body.Close())
+	return body, response.StatusCode
+}
+
+func verifySnapshot(t *testing.T, body []byte) []byte {
+	// reformat JSON body with indentation, to make diffing easier
+	var data interface{}
+	require.NoError(t, json.Unmarshal(body, &data))
+	body, err := json.MarshalIndent(data, "", "  ")
+	require.NoError(t, err)
+
+	snapshotFile := filepath.Join(snapshotLocation, strings.ReplaceAll(t.Name(), "/", "_")+".json")
+	if regenerateSnapshots {
+		os.WriteFile(snapshotFile, body, 0o644)
+	}
+	snapshot, err := os.ReadFile(snapshotFile)
+	require.NoError(t, err)
+	assert.Equal(t, string(snapshot), string(body), "comparing against stored snapshot. Use REGENERATE_SNAPSHOTS=true to rebuild snapshots.")
+	return body
+}
+
+func parseResponse(t *testing.T, body []byte, obj gogoproto.Message) {
+	require.NoError(t, gogojsonpb.Unmarshal(bytes.NewBuffer(body), obj))
+}
+
+func parseChunkResponse(t *testing.T, body []byte, obj gogoproto.Message) {
+	// Unwrap the 'result' container generated by the gateway.
+	// See https://github.com/grpc-ecosystem/grpc-gateway/issues/2189
+	type resultWrapper struct {
+		Result json.RawMessage `json:"result"`
+	}
+	var result resultWrapper
+	require.NoError(t, json.Unmarshal(body, &result))
+	parseResponse(t, result.Result, obj)
+}
+
+func makeTestTrace() (*model.Trace, model.TraceID) {
+	traceID := model.NewTraceID(150, 160)
+	return &model.Trace{
+		Spans: []*model.Span{
+			{
+				TraceID:       traceID,
+				SpanID:        model.NewSpanID(180),
+				OperationName: "foobar",
+			},
+		},
+	}, traceID
+}
+
 func testGRPCGateway(
 	t *testing.T, basePath string,
 	serverTLS, clientTLS *tlscfg.Options,
@@ -144,36 +222,94 @@ func testGRPCGatewayWithTenancy(
 	setupRequest func(*http.Request),
 ) {
 	gw := setupGRPCGateway(t, basePath, serverTLS, clientTLS, tenancyOptions)
+	t.Run("GetServices", func(t *testing.T) {
+		runGatewayGetServices(t, gw, setupRequest)
+	})
+	t.Run("GetOperations", func(t *testing.T) {
+		runGatewayGetOperations(t, gw, setupRequest)
+	})
+	t.Run("GetTrace", func(t *testing.T) {
+		runGatewayGetTrace(t, gw, setupRequest)
+	})
+	t.Run("FindTraces", func(t *testing.T) {
+		runGatewayFindTraces(t, gw, setupRequest)
+	})
+}
 
-	traceID := model.NewTraceID(150, 160)
-	gw.reader.On("GetTrace", matchContext, matchTraceID).Return(
-		&model.Trace{
-			Spans: []*model.Span{
-				{
-					TraceID:       traceID,
-					SpanID:        model.NewSpanID(180),
-					OperationName: "foobar",
-				},
-			},
-		}, nil).Once()
+func runGatewayGetServices(t *testing.T, gw *testGateway, setupRequest func(*http.Request)) {
+	gw.reader.On("GetServices", matchContext).Return([]string{"foo"}, nil).Once()
 
-	req, err := http.NewRequest(http.MethodGet, gw.url+"/api/v3/traces/123", nil)
-	require.NoError(t, err)
-	req.Header.Set("Content-Type", "application/json")
-	setupRequest(req)
-	response, err := http.DefaultClient.Do(req)
-	require.NoError(t, err)
-	body, err := io.ReadAll(response.Body)
-	require.NoError(t, err)
-	require.NoError(t, response.Body.Close())
+	body, statusCode := gw.execRequest(t, &gatewayRequest{
+		url:          "/api/v3/services",
+		setupRequest: setupRequest,
+	})
+	require.Equal(t, http.StatusOK, statusCode)
+	body = verifySnapshot(t, body)
+
+	var response api_v3.GetServicesResponse
+	parseResponse(t, body, &response)
+	assert.Equal(t, []string{"foo"}, response.Services)
+}
+
+func runGatewayGetOperations(t *testing.T, gw *testGateway, setupRequest func(*http.Request)) {
+	qp := spanstore.OperationQueryParameters{ServiceName: "foo", SpanKind: "server"}
+	gw.reader.
+		On("GetOperations", matchContext, qp).
+		Return([]spanstore.Operation{{Name: "get_users", SpanKind: "server"}}, nil).Once()
+
+	body, statusCode := gw.execRequest(t, &gatewayRequest{
+		url:          "/api/v3/operations?service=foo&span_kind=server",
+		setupRequest: setupRequest,
+	})
+	require.Equal(t, http.StatusOK, statusCode)
+	body = verifySnapshot(t, body)
+
+	var response api_v3.GetOperationsResponse
+	parseResponse(t, body, &response)
+	require.Len(t, response.Operations, 1)
+	assert.Equal(t, "get_users", response.Operations[0].Name)
+	assert.Equal(t, "server", response.Operations[0].SpanKind)
+}
+
+func runGatewayGetTrace(t *testing.T, gw *testGateway, setupRequest func(*http.Request)) {
+	trace, traceID := makeTestTrace()
+	gw.reader.On("GetTrace", matchContext, traceID).Return(trace, nil).Once()
+
+	body, statusCode := gw.execRequest(t, &gatewayRequest{
+		url:          "/api/v3/traces/" + traceID.String(), // hex string
+		setupRequest: setupRequest,
+	})
+	require.Equal(t, http.StatusOK, statusCode, "response=%s", string(body))
+	body = verifySnapshot(t, body)
 
-	jsonpb := &runtime.JSONPb{}
-	var envelope envelope
-	err = json.Unmarshal(body, &envelope)
-	require.NoError(t, err)
 	var spansResponse api_v3.SpansResponseChunk
-	err = jsonpb.Unmarshal(envelope.Result, &spansResponse)
-	require.NoError(t, err)
+	parseChunkResponse(t, body, &spansResponse)
+
+	assert.Len(t, spansResponse.GetResourceSpans(), 1)
+	assert.Equal(t, bytesOfTraceID(t, traceID.High, traceID.Low), spansResponse.GetResourceSpans()[0].GetScopeSpans()[0].GetSpans()[0].GetTraceId())
+}
+
+func runGatewayFindTraces(t *testing.T, gw *testGateway, setupRequest func(*http.Request)) {
+	trace, traceID := makeTestTrace()
+	gw.reader.
+		On("FindTraces", matchContext, mock.AnythingOfType("*spanstore.TraceQueryParameters")).
+		Return([]*model.Trace{trace}, nil).Once()
+
+	q := url.Values{}
+	q.Set("query.service_name", "foobar")
+	q.Set("query.start_time_min", time.Now().Format(time.RFC3339))
+	q.Set("query.start_time_max", time.Now().Format(time.RFC3339))
+
+	body, statusCode := gw.execRequest(t, &gatewayRequest{
+		url:          "/api/v3/traces?" + q.Encode(),
+		setupRequest: setupRequest,
+	})
+	require.Equal(t, http.StatusOK, statusCode, "response=%s", string(body))
+	body = verifySnapshot(t, body)
+
+	var spansResponse api_v3.SpansResponseChunk
+	parseChunkResponse(t, body, &spansResponse)
+
 	assert.Len(t, spansResponse.GetResourceSpans(), 1)
 	assert.Equal(t, bytesOfTraceID(t, traceID.High, traceID.Low), spansResponse.GetResourceSpans()[0].GetScopeSpans()[0].GetSpans()[0].GetTraceId())
 }
@@ -207,11 +343,6 @@ func TestGRPCGatewayWithBasePathAndTLS(t *testing.T) {
 	testGRPCGateway(t, "/jaeger", serverTLS, clientTLS)
 }
 
-// For more details why this is needed see https://github.com/grpc-ecosystem/grpc-gateway/issues/2189
-type envelope struct {
-	Result json.RawMessage `json:"result"`
-}
-
 func TestGRPCGatewayWithTenancy(t *testing.T) {
 	tenancyOptions := tenancy.Options{
 		Enabled: true,
@@ -226,7 +357,7 @@ func TestGRPCGatewayWithTenancy(t *testing.T) {
 		})
 }
 
-func TestTenancyGRPCRejection(t *testing.T) {
+func TestGRPCGatewayTenancyRejection(t *testing.T) {
 	basePath := "/"
 	tenancyOptions := tenancy.Options{Enabled: true}
 	gw := setupGRPCGateway(t,

cmd/query/app/apiv3/snapshots/TestGRPCGatewayWithBasePathAndTLS_FindTraces.json

@@ -0,0 +1,24 @@
+{
+  "result": {
+    "resourceSpans": [
+      {
+        "resource": {},
+        "scopeSpans": [
+          {
+            "scope": {},
+            "spans": [
+              {
+                "endTimeUnixNano": "11651379494838206464",
+                "name": "foobar",
+                "spanId": "AAAAAAAAALQ=",
+                "startTimeUnixNano": "11651379494838206464",
+                "status": {},
+                "traceId": "AAAAAAAAAJYAAAAAAAAAoA=="
+              }
+            ]
+          }
+        ]
+      }
+    ]
+  }
+}

cmd/query/app/apiv3/snapshots/TestGRPCGatewayWithBasePathAndTLS_GetOperations.json

@@ -0,0 +1,8 @@
+{
+  "operations": [
+    {
+      "name": "get_users",
+      "spanKind": "server"
+    }
+  ]
+}

cmd/query/app/apiv3/snapshots/TestGRPCGatewayWithBasePathAndTLS_GetServices.json

@@ -0,0 +1,5 @@
+{
+  "services": [
+    "foo"
+  ]
+}

cmd/query/app/apiv3/snapshots/TestGRPCGatewayWithBasePathAndTLS_GetTrace.json

@@ -0,0 +1,24 @@
+{
+  "result": {
+    "resourceSpans": [
+      {
+        "resource": {},
+        "scopeSpans": [
+          {
+            "scope": {},
+            "spans": [
+              {
+                "endTimeUnixNano": "11651379494838206464",
+                "name": "foobar",
+                "spanId": "AAAAAAAAALQ=",
+                "startTimeUnixNano": "11651379494838206464",
+                "status": {},
+                "traceId": "AAAAAAAAAJYAAAAAAAAAoA=="
+              }
+            ]
+          }
+        ]
+      }
+    ]
+  }
+}

cmd/query/app/apiv3/snapshots/TestGRPCGatewayWithTenancy_FindTraces.json

@@ -0,0 +1,24 @@
+{
+  "result": {
+    "resourceSpans": [
+      {
+        "resource": {},
+        "scopeSpans": [
+          {
+            "scope": {},
+            "spans": [
+              {
+                "endTimeUnixNano": "11651379494838206464",
+                "name": "foobar",
+                "spanId": "AAAAAAAAALQ=",
+                "startTimeUnixNano": "11651379494838206464",
+                "status": {},
+                "traceId": "AAAAAAAAAJYAAAAAAAAAoA=="
+              }
+            ]
+          }
+        ]
+      }
+    ]
+  }
+}

cmd/query/app/apiv3/snapshots/TestGRPCGatewayWithTenancy_GetOperations.json

@@ -0,0 +1,8 @@
+{
+  "operations": [
+    {
+      "name": "get_users",
+      "spanKind": "server"
+    }
+  ]
+}

cmd/query/app/apiv3/snapshots/TestGRPCGatewayWithTenancy_GetServices.json

@@ -0,0 +1,5 @@
+{
+  "services": [
+    "foo"
+  ]
+}