test: Add more tests (#5)

ayush6624 · Mar 18, 2023 · acfe8cf · acfe8cf
1 parent a222338
commit acfe8cf
Show file tree

Hide file tree

Showing 4 changed files with 138 additions and 46 deletions.
diff --git a/README.md b/README.md
@@ -28,7 +28,7 @@ ___
 1. In your Go code, import the ChatGPT-Go package:
     ```go
     import (
-        "github.com/<username>/chatgpt-go/chatgpt"
+        "github.com/ayush6624/chatgpt-go/chatgpt"
     )
     ```
 

diff --git a/chat.go b/chat.go
@@ -111,10 +111,7 @@ func (c *Client) Send(ctx context.Context, req *ChatCompletionRequest) (*ChatRes
 		return nil, err
 	}
 
-	reqBytes, err := json.Marshal(req)
-	if err != nil {
-		return nil, err
-	}
+	reqBytes, _ := json.Marshal(req)
 
 	endpoint := "/chat/completions"
 	httpReq, err := http.NewRequest("POST", c.config.BaseURL+endpoint, bytes.NewBuffer(reqBytes))

diff --git a/chat_test.go b/chat_test.go
@@ -1,10 +1,13 @@
 package chatgpt
 
 import (
+	"context"
+	"net/http"
+	"net/http/httptest"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
 	chatgpt_errors "github.com/ayush6624/go-chatgpt/utils"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestValidate(t *testing.T) {
@@ -56,9 +59,9 @@ func TestValidate(t *testing.T) {
 		{
 			name: "Invalid presence penalty",
 			request: &ChatCompletionRequest{
-				Model:            GPT35Turbo,
-				Messages:         validRequest().Messages,
-				PresencePenalty:  -3,
+				Model:           GPT35Turbo,
+				Messages:        validRequest().Messages,
+				PresencePenalty: -3,
 			},
 			expectedError: chatgpt_errors.ErrInvalidPresencePenalty,
 		},
@@ -92,3 +95,121 @@ func validRequest() *ChatCompletionRequest {
 		},
 	}
 }
+
+func newTestServerAndClient() (*httptest.Server, *Client) {
+	// Create a new test HTTP server to handle requests
+	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{ "id": "chatcmpl-abcd", "object": "chat.completion", "created_at": 0, "choices": [ { "index": 0, "message": { "role": "assistant", "content": "\n\n Sample response" }, "finish_reason": "stop" } ], "usage": { "prompt_tokens": 19, "completion_tokens": 47, "total_tokens": 66 }}`))
+	}))
+
+	// Create a new client with the test server's URL and a mock API key
+	return testServer, &Client{
+		client: http.DefaultClient,
+		config: &Config{
+			BaseURL:        testServer.URL,
+			APIKey:         "mock_api_key",
+			OrganizationID: "mock_organization_id",
+		},
+	}
+}
+
+func newTestClientWithInvalidResponse() (*httptest.Server, *Client) {
+	// Create a new test HTTP server to handle requests
+	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{ fakejson }`))
+	}))
+
+	// Create a new client with the test server's URL and a mock API key
+	return testServer, &Client{
+		client: http.DefaultClient,
+		config: &Config{
+			BaseURL:        testServer.URL,
+			APIKey:         "mock_api_key",
+			OrganizationID: "mock_organization_id",
+		},
+	}
+}
+
+func newTestClientWithInvalidStatusCode() (*httptest.Server, *Client) {
+	// Create a new test HTTP server to handle requests
+	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte(`{ "error": "bad request" }`))
+	}))
+
+	// Create a new client with the test server's URL and a mock API key
+	return testServer, &Client{
+		client: http.DefaultClient,
+		config: &Config{
+			BaseURL:        testServer.URL,
+			APIKey:         "mock_api_key",
+			OrganizationID: "mock_organization_id",
+		},
+	}
+}
+
+func TestSend(t *testing.T) {
+	server, client := newTestServerAndClient()
+	defer server.Close()
+
+	_, err := client.Send(context.Background(), &ChatCompletionRequest{
+		Model: GPT35Turbo,
+		Messages: []ChatMessage{
+			{
+				Role:    ChatGPTModelRoleUser,
+				Content: "Hello",
+			},
+		},
+	})
+	assert.NoError(t, err)
+
+	_, err = client.Send(context.Background(), &ChatCompletionRequest{
+		Model: "invalid model",
+		Messages: []ChatMessage{
+			{
+				Role:    ChatGPTModelRoleUser,
+				Content: "Hello",
+			},
+		},
+	})
+	assert.Error(t, err)
+
+	server, client = newTestClientWithInvalidResponse()
+	defer server.Close()
+
+	_, err = client.Send(context.Background(), &ChatCompletionRequest{
+		Model: GPT35Turbo,
+		Messages: []ChatMessage{
+			{
+				Role:    ChatGPTModelRoleUser,
+				Content: "Hello",
+			},
+		},
+	})
+	assert.Error(t, err)
+
+	server, client = newTestClientWithInvalidStatusCode()
+	defer server.Close()
+
+	_, err = client.Send(context.Background(), &ChatCompletionRequest{
+		Model: GPT35Turbo,
+		Messages: []ChatMessage{
+			{
+				Role:    ChatGPTModelRoleUser,
+				Content: "Hello",
+			},
+		},
+	})
+	assert.Error(t, err)
+
+}
+
+func TestSimpleSend(t *testing.T) {
+	server, client := newTestServerAndClient()
+	defer server.Close()
+
+	_, err := client.SimpleSend(context.Background(), "Hello")
+	assert.NoError(t, err)
+}
diff --git a/client_test.go b/client_test.go
@@ -4,10 +4,9 @@ import (
 	"context"
 	"fmt"
 	"net/http"
-	"net/http/httptest"
 	"testing"
 
-	"github.com/ayush6624/go-chatgpt/utils"
+	chatgpt_errors "github.com/ayush6624/go-chatgpt/utils"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -32,30 +31,17 @@ func TestNewClient(t *testing.T) {
 }
 
 func TestClient2_sendRequest(t *testing.T) {
-	// Create a new test HTTP server to handle requests
-	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.WriteHeader(http.StatusOK)
-	}))
-	defer testServer.Close()
-
-	// Create a new client with the test server's URL and a mock API key
-	client := &Client{
-		client: http.DefaultClient,
-		config: &Config{
-			BaseURL: testServer.URL,
-			APIKey:  "mock_api_key",
-			OrganizationID: "mock_organization_id",
-		},
-	}
+	// Create a new test HTTP server and client to handle requests
+	testServer, client := newTestServerAndClient()
 
 	// Create a new request
-	req, err := http.NewRequest("GET", testServer.URL, nil)
+	req, err := http.NewRequest("POST", testServer.URL, nil)
 	assert.NoError(t, err)
 
 	// Send the request using the ChatGPT client
-	resp, err := client.sendRequest(context.Background(), req)
+	res, err := client.sendRequest(context.Background(), req)
 	assert.NoError(t, err)
-	assert.Equal(t, http.StatusOK, resp.StatusCode)
+	assert.Equal(t, http.StatusOK, res.StatusCode)
 
 	expectedHeader := map[string]string{
 		"Authorization":       fmt.Sprintf("Bearer %s", "mock_api_key"),
@@ -64,29 +50,17 @@ func TestClient2_sendRequest(t *testing.T) {
 		"Accept":              "application/json",
 	}
 
-	// Check that the request's header is set correctly
+	// Check that the request's header is set correctly, after sendRequest was called
 	for key, value := range expectedHeader {
 		assert.Equal(t, req.Header.Get(key), value)
 	}
 
-	tServer2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.WriteHeader(http.StatusBadRequest)
-		w.Write([]byte(`{"result": "failure"}`))
-	}))
-	defer tServer2.Close()
-
-	client, err = NewClientWithConfig(&Config{
-		BaseURL: tServer2.URL,
-		APIKey:  "mock_api_key",
-		OrganizationID: "mock_organization_id",
-	})
-	assert.NoError(t, err)
-
+	testServer, client = newTestClientWithInvalidStatusCode()
 	// Prepare a test request
-	req, err = http.NewRequest("GET", tServer2.URL, nil)
+	req, err = http.NewRequest("GET", testServer.URL, nil)
 	assert.NoError(t, err)
 
-	// Send the request using the ChatGPT client
-	resp, err = client.sendRequest(context.Background(), req)
+	resp, err := client.sendRequest(context.Background(), req)
 	assert.Error(t, err)
+	assert.Nil(t, resp)
 }