test(cody): Add unit tests for the Completions API (#63434)

The [Server-side Cody Model
Selection](https://linear.app/sourcegraph/project/server-side-cody-model-selection-cca47c48da6d)
project requires refactoring a lot of how the Completions API endpoint
works for the Sourcegraph backend. It will need to update any place
where we interact with the site config as it relates to LLMs, or any
time we figure out which models can-or-can-not be used.

In order to safely land those refactoring without breaking existing
users, who will be using the "older config format", we need tests. LOTS
and LOTS of tests.

This PR adds the necessary infrastructure for writing unit tests against
the Completions API, and adds a few basic ones for the Anthropic API
provider. I wouldn't say it's as easy as we'd like to write these tests,
but at least now it is _possible_. And we can further streamline things
from here.

## Overview

The bulk of functionality is in `entconfig_base_test.go`. That file
contains some basic unit tests for how the enterprise configuration data
is loaded, and provides the mocks and test infrastructure.

The crux of which is the `apiProviderTestInfra` struct. It will bundle
mocks and the HTTP handlers to test, and provides high-level methods for
invoking the completion API.

```go
type apiProviderTestInfra struct {}

func (ti *apiProviderTestInfra) PushGetModelResult(model string, err error)
func (ti *apiProviderTestInfra) SetSiteConfig(siteConfig schema.SiteConfiguration)
func (ti *apiProviderTestInfra) CallChatCompletionAPI(t *testing.T, reqObj types.CodyCompletionRequestParameters) (int, string)
func (ti *apiProviderTestInfra) CallCodeCompletionAPI(t *testing.T, reqObj types.CodyCompletionRequestParameters) (int, string)
func (ti *apiProviderTestInfra) AssertCompletionsResponse(t *testing.T, rawResponseJSON string, wantResponse types.CompletionResponse)
```

What gets trick however, is how we actually hook into the implementation
details of the completions API endpoint. You see, the user makes an HTTP
request to the Sourcegraph instance. But we then figure out which LLM
model to use, and then build an HTTP request to send to the specific LLM
API Provider.

So the test infrastructure allows you to mock out that "middle part".
The
`AssertCodyGatewayReceivesRequestWithResponse`/`AssertGenericExternalAPIRequestWithResponse`
functions will:

1. Verify that the Sourcegraph instance is making an API call to the LLM
provider that matches the format we expect. (e.g. using the correct
Anthropic API request, etc.)
2. The outbound HTTP request looks like it should, that it contains the
right authorization headers, URL path, etc. (e.g. is it using the API
key from the site config?)
3. Finally, it returns the HTTP response from the API provider. (i.e.
whatever Anthropic or Cody Gateway would have returned.)

```go
type assertLLMRequestOptions struct {
	// WantRequestObj is what we expect the outbound HTTP request's JSON body
	// to be equal to. Required.
	WantRequestObj any
	// OutResponseObj is serialized to JSON and sent to the caller, i.e. our
	// LLM API Provider which is making the API request. Required.
	OutResponseObj any

	// WantRequestPath is the URL Path we expect in the outbound HTTP request.
	// No check is done if empty.
	WantRequestPath string

	// WantHeaders are HTTP header key/value pairs that must be present.
	WantHeaders map[string]string
}
func (ti *apiProviderTestInfra) AssertCodyGatewayReceivesRequestWithResponse(
	t *testing.T, opts assertLLMRequestOptions)
func (ti *apiProviderTestInfra) AssertGenericExternalAPIRequestWithResponse(
	t *testing.T, opts assertLLMRequestOptions)
```

Unfortunately it's super gnarly because we aren't really exposing the
API data types from the API providers in a useful way. (e.g. perhaps we
should just use the standard Anthropic Golang API client library.) So
generating the test data relies on a lot of `map[string]any` to make it
easier to construct arbitrary data types that can serialize to JSON the
way that we need them to.

Anyways, with all of this test infrastructure in-place, you can write
API provider tests like the following. Here's one such test from
`entconfig_anthropic_test.go` which confirms that various aspects of the
site-configuration are honored correctly when configured to use BYOK
mode.

I've added "ℹ️ comments" to highlight some of the tricker parts.


```go
t.Run("ViaBYOK", func(t *testing.T) {
		const (
			anthropicAPIKeyInConfig      = "secret-api-key"
			anthropicAPIEndpointInConfig = "https://byok.anthropic.com/path/from/config"
			chatModelInConfig            = "anthropic/claude-3-opus"
			codeModelInConfig            = "anthropic/claude-3-haiku"
		)

		infra.SetSiteConfig(schema.SiteConfiguration{
			CodyEnabled:                  pointers.Ptr(true),
			CodyPermissions:              pointers.Ptr(false),
			CodyRestrictUsersFeatureFlag: pointers.Ptr(false),

			// LicenseKey is required in order to use Cody.
			LicenseKey: "license-key",
			Completions: &schema.Completions{
				Provider:    "anthropic",
				AccessToken: anthropicAPIKeyInConfig,
				Endpoint:    anthropicAPIEndpointInConfig,

				ChatModel:       chatModelInConfig,
				CompletionModel: codeModelInConfig,
			},
		})

		t.Run("ChatModel", func(t *testing.T) {
			// ℹ️ Generating the "wantAnthropicRequest" and "outAnthropicResponse"
			// data is super-tedious. So we instead have a single function
			// that returns "a valid set", that we then customize.
			// So here, we just update the model we expect to see in the API
			// call to Anthropic.

			// Start with the stock test data, but customize it to reflect
			// what we expect to see based on the site configuration.
			testData := getValidTestData()
			testData.OutboundAnthropicRequest["model"] = "anthropic/claude-3-opus"

			// Register our hook to verify Cody Gateway got called with
			// the requested data.
			infra.AssertGenericExternalAPIRequestWithResponse(
				t, assertLLMRequestOptions{
					WantRequestPath: "/path/from/config",
					WantRequestObj:  &testData.OutboundAnthropicRequest,
					OutResponseObj:  &testData.InboundAnthropicRequest,
					WantHeaders: map[string]string{
						// Yes, Anthropic's API uses "X-Api-Key" rather than the "Authorization" header. 🤷
						"X-Api-Key": anthropicAPIKeyInConfig,
					},
				})

			// ℹ️ This `PushGetModelResult` is just a quirk of how the code
			// under test works. We mock out the `getModelFn` that is invoked
			// to resolve the _actual_ LLM model to use. (And not necessarily
			// use the one from the HTTP request.)
			infra.PushGetModelResult(chatModelInConfig, nil)

			status, responseBody := infra.CallChatCompletionAPI(t, testData.InitialCompletionRequest)

			assert.Equal(t, http.StatusOK, status)
			infra.AssertCompletionsResponse(t, responseBody, types.CompletionResponse{
				// ℹ️ The "totally rewrite it in Rust!" is coming from the
				// fake Anthropic response, from `getValidTestData`.
				Completion: "you should totally rewrite it in Rust!",
				StopReason: "max_tokens",
				Logprobs:   nil,
			})
		})
	})
```

## Next steps

Once this gets checked-in, my plan is to carefully add new unit tests
for the existing functionality before DISMANTLING the code to write
through an entirely different site configuration object. 🤞 this will
allow me to do so in such a way that I can confirm my changes won't
alter any existing Sourcegraph instances that are using the older
configuration format.

## Test plan

Adds more tests.

## Changelog

NA. Just trivial changes and adding more tests.

Author: chrsmith

cmd/frontend/internal/httpapi/completions/BUILD.bazel

@@ -51,15 +51,35 @@ go_library(
 
 go_test(
     name = "completions_test",
-    srcs = ["handler_test.go"],
+    srcs = [
+        "entconfig_anthropic_test.go",
+        "entconfig_base_test.go",
+        "get_model_test.go",
+        "handler_test.go",
+    ],
     embed = [":completions"],
-    tags = [TAG_CODY_CORE],
+    tags = [
+        TAG_CODY_CORE,
+        # Test indirectly localhost Redis.
+        "requires-network",
+    ],
     deps = [
+        "//internal/actor",
         "//internal/completions/types",
         "//internal/conf",
+        "//internal/conf/conftypes",
         "//internal/database/dbmocks",
         "//internal/featureflag",
+        "//internal/httpcli",
+        "//internal/licensing",
+        "//internal/rcache",
+        "//internal/telemetry",
+        "//internal/telemetry/telemetrytest",
+        "//lib/errors",
+        "//lib/pointers",
         "//schema",
+        "@com_github_sourcegraph_log//logtest",
+        "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
     ],
 )

cmd/frontend/internal/httpapi/completions/entconfig_anthropic_test.go

@@ -0,0 +1,188 @@
+package completions
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/sourcegraph/sourcegraph/internal/completions/types"
+	"github.com/sourcegraph/sourcegraph/lib/pointers"
+	"github.com/sourcegraph/sourcegraph/schema"
+)
+
+func testAPIProviderAnthropic(t *testing.T, infra *apiProviderTestInfra) {
+	// validAnthropicRequestData bundles the messages sent between major
+	// components of the Completions API.
+	type validAnthropicRequestData struct {
+		// InitialCompletionRequest is the request sent by the user to the
+		// Sourcegraph instance.
+		InitialCompletionRequest types.CodyCompletionRequestParameters
+
+		// OutboundAnthropicRequest is the data sent from this Sourcegraph
+		// instance to the API Provider (e.g. Anthropic, Cody Gateway, AWS
+		// Bedrock, etc.)
+		OutboundAnthropicRequest map[string]any
+
+		// InboundAnthropicRequest is the response from the API Provider
+		// to the Sourcegraph instance.
+		InboundAnthropicRequest map[string]any
+	}
+
+	// getValidTestData returns a valid set of request data.
+	getValidTestData := func() validAnthropicRequestData {
+		initialCompletionRequest := types.CodyCompletionRequestParameters{
+			CompletionRequestParameters: types.CompletionRequestParameters{
+				Messages: []types.Message{
+					{
+						Speaker: "human",
+						Text:    "please make this code better",
+					},
+				},
+				Stream: pointers.Ptr(false),
+			},
+		}
+
+		// Anthropic-specific request object we expect to see sent to Cody Gateway.
+		// See `anthropicRequestParameters`.
+		outboundAnthropicRequest := map[string]any{
+			"model": "claude-2.0",
+			"messages": []map[string]any{
+				{
+					"role": "user",
+					"content": []map[string]any{
+						{
+							"type": "text",
+							"text": "please make this code better",
+						},
+					},
+				},
+			},
+		}
+
+		// Stock response we would receive from Anthropic.
+		//
+		// The expected output is found also defined in the Anthropic completion provider codebase,
+		// as `anthropicNonStreamingResponse`.` But it's easier to keep those types unexported.
+		inboundAnthropicRequest := map[string]any{
+			"content": []map[string]string{
+				{
+					"speak": "user",
+					"text":  "you should totally rewrite it in Rust!",
+				},
+			},
+			"usage": map[string]int{
+				"input_token":   100,
+				"output_tokens": 200,
+			},
+			"stop_reason": "max_tokens",
+		}
+
+		return validAnthropicRequestData{
+			InitialCompletionRequest: initialCompletionRequest,
+			OutboundAnthropicRequest: outboundAnthropicRequest,
+			InboundAnthropicRequest:  inboundAnthropicRequest,
+		}
+	}
+
+	t.Run("WithDefaultConfig", func(t *testing.T) {
+		infra.SetSiteConfig(schema.SiteConfiguration{
+			CodyEnabled:                  pointers.Ptr(true),
+			CodyPermissions:              pointers.Ptr(false),
+			CodyRestrictUsersFeatureFlag: pointers.Ptr(false),
+
+			// LicenseKey is required in order to use Cody, but other than
+			// that we don't provide any "completions" configuration.
+			// This will default to Anthropic models.
+			LicenseKey:  "license-key",
+			Completions: nil,
+		})
+
+		// Confirm that the default configuration `Completions: nil` will use
+		// Cody Gateway as the LLM API Provider for the Anthropic models.
+		t.Run("ViaCodyGateway", func(t *testing.T) {
+			// The Model isn't included in the CompletionRequestParameters, so we have the getModelFn callback
+			// return claude-2. The Site Configuration will then route this to Cody Gateway (and not BYOK Anthropic),
+			// and we sanity check the request to Cody Gateway matches what is expected, and we serve a valid response.
+			infra.PushGetModelResult("anthropic/claude-2", nil)
+
+			// Generate some basic test data and confirm that the completions handler
+			// code works as expected.
+			testData := getValidTestData()
+
+			// Register our hook to verify Cody Gateway got called with
+			// the requested data.
+			infra.AssertCodyGatewayReceivesRequestWithResponse(
+				t, assertLLMRequestOptions{
+					WantRequestPath: "/v1/completions/anthropic-messages",
+					WantRequestObj:  &testData.OutboundAnthropicRequest,
+					OutResponseObj:  &testData.InboundAnthropicRequest,
+				})
+
+			status, responseBody := infra.CallChatCompletionAPI(t, testData.InitialCompletionRequest)
+
+			assert.Equal(t, http.StatusOK, status)
+			infra.AssertCompletionsResponse(t, responseBody, types.CompletionResponse{
+				Completion: "you should totally rewrite it in Rust!",
+				StopReason: "max_tokens",
+				Logprobs:   nil,
+			})
+		})
+	})
+
+	t.Run("ViaBYOK", func(t *testing.T) {
+		const (
+			anthropicAPIKeyInConfig      = "secret-api-key"
+			anthropicAPIEndpointInConfig = "https://byok.anthropic.com/path/from/config"
+			chatModelInConfig            = "anthropic/claude-3-opus"
+			codeModelInConfig            = "anthropic/claude-3-haiku"
+		)
+
+		infra.SetSiteConfig(schema.SiteConfiguration{
+			CodyEnabled:                  pointers.Ptr(true),
+			CodyPermissions:              pointers.Ptr(false),
+			CodyRestrictUsersFeatureFlag: pointers.Ptr(false),
+
+			// LicenseKey is required in order to use Cody.
+			LicenseKey: "license-key",
+			Completions: &schema.Completions{
+				Provider:    "anthropic",
+				AccessToken: anthropicAPIKeyInConfig,
+				Endpoint:    anthropicAPIEndpointInConfig,
+
+				ChatModel:       chatModelInConfig,
+				CompletionModel: codeModelInConfig,
+			},
+		})
+
+		t.Run("ChatModel", func(t *testing.T) {
+			// Start with the stock test data, but customize it to reflect
+			// what we expect to see based on the site configuration.
+			testData := getValidTestData()
+			testData.OutboundAnthropicRequest["model"] = "anthropic/claude-3-opus"
+
+			// Register our hook to verify Cody Gateway got called with
+			// the requested data.
+			infra.AssertGenericExternalAPIRequestWithResponse(
+				t, assertLLMRequestOptions{
+					WantRequestPath: "/path/from/config",
+					WantRequestObj:  &testData.OutboundAnthropicRequest,
+					OutResponseObj:  &testData.InboundAnthropicRequest,
+					WantHeaders: map[string]string{
+						// Yes, Anthropic's API uses "X-Api-Key" rather than the "Authorization" header. 🤷
+						"X-Api-Key": anthropicAPIKeyInConfig,
+					},
+				})
+
+			infra.PushGetModelResult(chatModelInConfig, nil)
+			status, responseBody := infra.CallChatCompletionAPI(t, testData.InitialCompletionRequest)
+
+			assert.Equal(t, http.StatusOK, status)
+			infra.AssertCompletionsResponse(t, responseBody, types.CompletionResponse{
+				Completion: "you should totally rewrite it in Rust!",
+				StopReason: "max_tokens",
+				Logprobs:   nil,
+			})
+		})
+	})
+}