From 646989cc5bb61f73335017243e0b008b149ba0ab Mon Sep 17 00:00:00 2001
From: Rich Coggins <coggsfl@gmail.com>
Date: Wed, 14 Jun 2023 10:19:18 -0400
Subject: [PATCH] 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 <richard.coggins@officedepot.com>
---
 chat_test.go            | 18 ++++++++++++++++++
 internal/test/server.go | 14 +++++++++-----
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/chat_test.go b/chat_test.go
index ebe29f9eb..a43bb4aa6 100644
--- a/chat_test.go
+++ b/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
diff --git a/internal/test/server.go b/internal/test/server.go
index 79d55c405..3813ff869 100644
--- a/internal/test/server.go
+++ b/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)
 	}))
 }