From 6541aa5f6e0a845cc96232786173f786e8774322 Mon Sep 17 00:00:00 2001
From: Ed Snible <snible@us.ibm.com>
Date: Mon, 11 Jul 2022 14:38:11 -0400
Subject: [PATCH] HTTP tenancy test

Signed-off-by: Ed Snible <snible@us.ibm.com>
---
 cmd/query/app/grpc_handler_test.go |  1 -
 cmd/query/app/http_handler_test.go | 43 ++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/cmd/query/app/grpc_handler_test.go b/cmd/query/app/grpc_handler_test.go
index bc8b40f4905..e28a81da095 100644
--- a/cmd/query/app/grpc_handler_test.go
+++ b/cmd/query/app/grpc_handler_test.go
@@ -1182,7 +1182,6 @@ func TestTenancyContextFlowGRPC(t *testing.T) {
 
 		for tenant, expected := range allExpectedResults {
 			t.Run(tenant, func(t *testing.T) {
-
 				// Test context propagation to Unary method.
 				resGetServices, err := client.GetServices(withOutgoingMetadata(t, context.Background(), tc.Header, tenant), &api_v2.GetServicesRequest{})
 				require.NoError(t, err, "expecting gRPC to succeed with %q tenancy header", tenant)
diff --git a/cmd/query/app/http_handler_test.go b/cmd/query/app/http_handler_test.go
index 0238dbfc70a..2ad1cb7094b 100644
--- a/cmd/query/app/http_handler_test.go
+++ b/cmd/query/app/http_handler_test.go
@@ -17,6 +17,7 @@ package app
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -933,3 +934,45 @@ func TestSearchTenancyRejectionHTTP(t *testing.T) {
 	require.Equal(t, http.StatusOK, resp.StatusCode)
 	// Skip unmarshal of response; it is enough that it succeeded
 }
+
+func TestSearchTenancyFlowTenantHTTP(t *testing.T) {
+	tenancyOptions := tenancy.Options{
+		Enabled: true,
+	}
+	ts := initializeTestServerWithOptions(
+		*tenancy.NewTenancyConfig(&tenancyOptions),
+		querysvc.QueryServiceOptions{})
+	defer ts.server.Close()
+	ts.spanReader.On("GetTrace", mock.MatchedBy(func(v interface{}) bool {
+		ctx, ok := v.(context.Context)
+		if !ok || tenancy.GetTenant(ctx) != "acme" {
+			return false
+		}
+		return true
+	}), mock.AnythingOfType("model.TraceID")).Return(mockTrace, nil).Twice()
+	ts.spanReader.On("GetTrace", mock.MatchedBy(func(v interface{}) bool {
+		ctx, ok := v.(context.Context)
+		if !ok || tenancy.GetTenant(ctx) != "megacorp" {
+			return false
+		}
+		return true
+	}), mock.AnythingOfType("model.TraceID")).Return(nil, errStorage).Once()
+
+	var responseAcme structuredResponse
+	err := getJSONCustomHeaders(
+		ts.server.URL+`/api/traces?traceID=1&traceID=2`,
+		map[string]string{"x-tenant": "acme"},
+		&responseAcme)
+	assert.NoError(t, err)
+	assert.Len(t, responseAcme.Errors, 0)
+	assert.Len(t, responseAcme.Data, 2)
+
+	var responseMegacorp structuredResponse
+	err = getJSONCustomHeaders(
+		ts.server.URL+`/api/traces?traceID=1&traceID=2`,
+		map[string]string{"x-tenant": "megacorp"},
+		&responseMegacorp)
+	assert.Contains(t, err.Error(), "storage error")
+	assert.Len(t, responseMegacorp.Errors, 0)
+	assert.Nil(t, responseMegacorp.Data)
+}