one-api/common/test/server.go

package test

import (
	"log"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strings"
)

const testAPI = "this-is-my-secure-token-do-not-steal!!"

func GetTestToken() string {
	return testAPI
}

type ServerTest struct {
	handlers map[string]handler
}
type handler func(w http.ResponseWriter, r *http.Request)

func NewTestServer() *ServerTest {
	return &ServerTest{handlers: make(map[string]handler)}
}

func OpenAICheck(w http.ResponseWriter, r *http.Request) bool {
	if r.Header.Get("Authorization") != "Bearer "+GetTestToken() && r.Header.Get("api-key") != GetTestToken() {
		w.WriteHeader(http.StatusUnauthorized)
		return false
	}
	return true
}

func (ts *ServerTest) RegisterHandler(path string, handler handler) {
	// to make the registered paths friendlier to a regex match in the route handler
	// in OpenAITestServer
	path = strings.ReplaceAll(path, "*", ".*")
	ts.handlers[path] = handler
}

// OpenAITestServer Creates a mocked OpenAI server which can pretend to handle requests during testing.
func (ts *ServerTest) TestServer(headerCheck func(w http.ResponseWriter, r *http.Request) bool) *httptest.Server {
	return httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("received a %s request at path %q\n", r.Method, r.URL.Path)

		// check auth
		if headerCheck != nil && !headerCheck(w, r) {
			return
		}

		// Handle /path/* routes.
		// Note: the * is converted to a .* in register handler for proper regex handling
		for route, handler := range ts.handlers {
			// Adding ^ and $ to make path matching deterministic since go map iteration isn't ordered
			pattern, _ := regexp.Compile("^" + route + "$")
			if pattern.MatchString(r.URL.Path) {
				handler(w, r)
				return
			}
		}
		http.Error(w, "the resource path doesn't exist", http.StatusNotFound)
	}))
}