Improve (#356) to support registration of wildcard URLs (#359)

* Improve (#356) to support registration of wildcard URLs

* Add TestAzureChatCompletions & TestAzureChatCompletionsWithCustomDeploymentName

* Remove TestAzureChatCompletionsWithCustomDeploymentName

---------

Co-authored-by: coggsflod <[email protected]>

Author: coggsfl

chat_test.go

@@ -67,6 +67,24 @@ func TestChatCompletions(t *testing.T) {
 	checks.NoError(t, err, "CreateChatCompletion error")
 }
 
+func TestAzureChatCompletions(t *testing.T) {
+	client, server, teardown := setupAzureTestServer()
+	defer teardown()
+	server.RegisterHandler("/openai/deployments/*", handleChatCompletionEndpoint)
+
+	_, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
+		MaxTokens: 5,
+		Model:     GPT3Dot5Turbo,
+		Messages: []ChatCompletionMessage{
+			{
+				Role:    ChatMessageRoleUser,
+				Content: "Hello!",
+			},
+		},
+	})
+	checks.NoError(t, err, "CreateAzureChatCompletion error")
+}
+
 // handleChatCompletionEndpoint Handles the ChatGPT completion endpoint by the test server.
 func handleChatCompletionEndpoint(w http.ResponseWriter, r *http.Request) {
 	var err error

internal/test/server.go

@@ -4,6 +4,7 @@ import (
 	"log"
 	"net/http"
 	"net/http/httptest"
+	"regexp"
 )
 
 const testAPI = "this-is-my-secure-token-do-not-steal!!"
@@ -36,11 +37,14 @@ func (ts *ServerTest) OpenAITestServer() *httptest.Server {
 			return
 		}
 
-		handlerCall, ok := ts.handlers[r.URL.Path]
-		if !ok {
-			http.Error(w, "the resource path doesn't exist", http.StatusNotFound)
-			return
+		// Handle /path/* routes.
+		for route, handler := range ts.handlers {
+			pattern, _ := regexp.Compile(route)
+			if pattern.MatchString(r.URL.Path) {
+				handler(w, r)
+				return
+			}
 		}
-		handlerCall(w, r)
+		http.Error(w, "the resource path doesn't exist", http.StatusNotFound)
 	}))
 }