more use of testify

liushuangls · Sep 3, 2024 · e8bfe81 · e8bfe81
1 parent d77c1b0
commit e8bfe81
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 30 deletions.
diff --git a/client_test.go b/client_test.go
@@ -3,32 +3,29 @@ package anthropic
 import (
 	"net/http"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestWithBetaVersion(t *testing.T) {
+	is := require.New(t)
 	t.Run("single beta version", func(t *testing.T) {
 		opt := withBetaVersion("fake-version")
 		request, err := http.NewRequest("GET", "http://example.com", nil)
-		if err != nil {
-			t.Fatalf("http.NewRequest error: %s", err)
-		}
+		is.NoError(err)
+
 		opt(request)
 
-		if req := request.Header.Get("anthropic-beta"); req != "fake-version" {
-			t.Errorf("unexpected BetaVersion: %s", req)
-		}
+		is.Equal("fake-version", request.Header.Get("anthropic-beta"))
 	})
 
 	t.Run("multiple beta versions", func(t *testing.T) {
 		opt := withBetaVersion("fake-version1", "fake-version2")
 		request, err := http.NewRequest("GET", "http://example.com", nil)
-		if err != nil {
-			t.Fatalf("http.NewRequest error: %s", err)
-		}
+		is.NoError(err)
+
 		opt(request)
 
-		if req := request.Header.Get("anthropic-beta"); req != "fake-version1,fake-version2" {
-			t.Errorf("unexpected BetaVersion: %s", req)
-		}
+		is.Equal("fake-version1,fake-version2", request.Header.Get("anthropic-beta"))
 	})
 }
diff --git a/config_test.go b/config_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/liushuangls/go-anthropic/v2"
+	"github.com/stretchr/testify/require"
 )
 
 func TestWithBaseURL(t *testing.T) {
@@ -20,20 +21,20 @@ func TestWithBaseURL(t *testing.T) {
 }
 
 func TestWithAPIVersion(t *testing.T) {
+	is := require.New(t)
 	t.Run("single generic version", func(t *testing.T) {
 		fakeAPIVersion := anthropic.APIVersion("fake-version")
 		opt := anthropic.WithAPIVersion(fakeAPIVersion)
 
 		c := anthropic.ClientConfig{}
 		opt(&c)
 
-		if c.APIVersion != fakeAPIVersion {
-			t.Errorf("unexpected APIVersion: %s", c.APIVersion)
-		}
+		is.Equal(fakeAPIVersion, c.APIVersion)
 	})
 }
 
 func TestWithHTTPClient(t *testing.T) {
+	is := require.New(t)
 	fakeHTTPClient := http.Client{}
 	fakeHTTPClient.Timeout = 1234
 
@@ -42,45 +43,39 @@ func TestWithHTTPClient(t *testing.T) {
 	c := anthropic.ClientConfig{}
 	opt(&c)
 
-	if c.HTTPClient != &fakeHTTPClient {
-		t.Errorf("unexpected HTTPClient: %v", c.HTTPClient)
-	}
+	is.Equal(&fakeHTTPClient, c.HTTPClient)
 }
 
 func TestWithEmptyMessagesLimit(t *testing.T) {
+	is := require.New(t)
 	fakeLimit := uint(1234)
 	opt := anthropic.WithEmptyMessagesLimit(fakeLimit)
 
 	c := anthropic.ClientConfig{}
 	opt(&c)
 
-	if c.EmptyMessagesLimit != fakeLimit {
-		t.Errorf("unexpected EmptyMessagesLimit: %d", c.EmptyMessagesLimit)
-	}
+	is.Equal(fakeLimit, c.EmptyMessagesLimit)
 }
 
 func TestWithBetaVersion(t *testing.T) {
+	is := require.New(t)
 	t.Run("single generic version", func(t *testing.T) {
 		fakeBetaVersion := anthropic.BetaVersion("fake-version")
 		opt := anthropic.WithBetaVersion(fakeBetaVersion)
 		c := anthropic.ClientConfig{}
 		opt(&c)
 
-		if c.BetaVersion[0] != fakeBetaVersion {
-			t.Errorf("unexpected BetaVersion: %s", c.BetaVersion)
-		}
-		if len(c.BetaVersion) != 1 {
-			t.Errorf("unexpected BetaVersion length: %d", len(c.BetaVersion))
-		}
+		is.Equal(fakeBetaVersion, c.BetaVersion[0])
+		is.Equal(1, len(c.BetaVersion))
 	})
 
 	t.Run("multiple versions", func(t *testing.T) {
 		opt := anthropic.WithBetaVersion("foo", "bar")
 		c := anthropic.ClientConfig{}
 		opt(&c)
 
-		if len(c.BetaVersion) != 2 {
-			t.Errorf("unexpected BetaVersion length: %d", len(c.BetaVersion))
-		}
+		is.Equal(2, len(c.BetaVersion))
+		is.Equal(anthropic.BetaVersion("foo"), c.BetaVersion[0])
+		is.Equal(anthropic.BetaVersion("bar"), c.BetaVersion[1])
 	})
 }
diff --git a/message_stream_test.go b/message_stream_test.go
@@ -104,6 +104,7 @@ func TestMessagesStreamError(t *testing.T) {
 }
 
 func TestCreateMessagesStream(t *testing.T) {
+	is := require.New(t)
 	t.Run("Does not error for empty unknown messages below limit", func(t *testing.T) {
 		emptyMessagesLimit := 100
 		server := test.NewTestServer()