Split tests into separate files (#22)

cohere-ai · May 13, 2022 · 618a95f · 618a95f
1 parent 1768ae1
commit 618a95f
Show file tree

Hide file tree

Showing 5 changed files with 227 additions and 207 deletions.
diff --git a/classify_test.go b/classify_test.go
@@ -0,0 +1,87 @@
+package cohere
+
+import (
+	"testing"
+)
+
+func TestClassify(t *testing.T) {
+	co, err := CreateClient(apiKey)
+	if err != nil {
+		t.Error(err)
+	}
+
+	t.Run("ClassifySuccessMinimumFields", func(t *testing.T) {
+		res, err := co.Classify("medium", ClassifyOptions{
+			Inputs: []string{"purple"},
+			Examples: []Example{
+				{"apple", "fruit"}, {"banana", "fruit"}, {"watermelon", "fruit"}, {"cherry", "fruit"}, {"lemon", "fruit"},
+				{"red", "color"}, {"blue", "color"}, {"blue", "color"}, {"yellow", "color"}, {"green", "color"}},
+		})
+
+		if err != nil {
+			t.Errorf("Expected result, got error: %s", err.Error())
+		}
+
+		if res.Classifications[0].Prediction != "color" {
+			t.Errorf("Expected: color. Receieved: %s", res.Classifications[0].Prediction)
+		}
+	})
+
+	t.Run("ClassifySuccessAllFields", func(t *testing.T) {
+		res, err := co.Classify("medium", ClassifyOptions{
+			TaskDescription: "Classify these words as either a color or a fruit.",
+			Inputs:          []string{"grape", "pink"},
+			Examples: []Example{
+				{"apple", "fruit"}, {"banana", "fruit"}, {"watermelon", "fruit"}, {"cherry", "fruit"}, {"lemon", "fruit"},
+				{"red", "color"}, {"blue", "color"}, {"blue", "color"}, {"yellow", "color"}, {"green", "color"}},
+			OutputIndicator: "This is a",
+		})
+
+		if err != nil {
+			t.Errorf("Expected result, got error: %s", err.Error())
+		}
+
+		if res.Classifications[0].Prediction != "fruit" {
+			t.Errorf("Expected: fruit. Receieved: %s", res.Classifications[0].Prediction)
+		}
+		if res.Classifications[1].Prediction != "color" {
+			t.Errorf("Expected: color. Receieved: %s", res.Classifications[1].Prediction)
+		}
+	})
+
+	t.Run("ClassifySuccessTaskDescription", func(t *testing.T) {
+		res, err := co.Classify("medium", ClassifyOptions{
+			TaskDescription: "Classify these words as a fruit or a color",
+			Inputs:          []string{"kiwi"},
+			Examples: []Example{
+				{"apple", "fruit"}, {"banana", "fruit"}, {"watermelon", "fruit"}, {"cherry", "fruit"}, {"lemon", "fruit"},
+				{"red", "color"}, {"blue", "color"}, {"blue", "color"}, {"yellow", "color"}, {"green", "color"}},
+		})
+
+		if err != nil {
+			t.Errorf("Expected result, got error: %s", err.Error())
+		}
+
+		if res.Classifications[0].Prediction != "fruit" {
+			t.Errorf("Expected: fruit. Receieved: %s", res.Classifications[0].Prediction)
+		}
+	})
+
+	t.Run("ClassifySuccessOutputIndicator", func(t *testing.T) {
+		res, err := co.Classify("medium", ClassifyOptions{
+			Inputs: []string{"pineapple"},
+			Examples: []Example{
+				{"apple", "fruit"}, {"banana", "fruit"}, {"watermelon", "fruit"}, {"cherry", "fruit"}, {"lemon", "fruit"},
+				{"red", "color"}, {"blue", "color"}, {"blue", "color"}, {"yellow", "color"}, {"green", "color"}},
+			OutputIndicator: "This is a",
+		})
+
+		if err != nil {
+			t.Errorf("Expected result, got error: %s", err.Error())
+		}
+
+		if res.Classifications[0].Prediction != "fruit" {
+			t.Errorf("Expected: fruit. Receieved: %s", res.Classifications[0].Prediction)
+		}
+	})
+}
diff --git a/client_test.go b/client_test.go
@@ -26,210 +26,3 @@ func TestErrors(t *testing.T) {
 		}
 	})
 }
-
-func TestGenerate(t *testing.T) {
-	co, err := CreateClient(apiKey)
-	if err != nil {
-		t.Error(err)
-	}
-
-	t.Run("Generate basic", func(t *testing.T) {
-		_, err := co.Generate("medium", GenerateOptions{
-			Prompt:      "Hello my name is",
-			MaxTokens:   10,
-			Temperature: 0.75,
-		})
-		if err != nil {
-			t.Errorf("expected result, got error: %s", err.Error())
-		}
-	})
-
-	t.Run("Generate multi", func(t *testing.T) {
-		num := 4
-		res, err := co.Generate("medium", GenerateOptions{
-			Prompt:         "What is your",
-			MaxTokens:      10,
-			Temperature:    0.75,
-			NumGenerations: num,
-		})
-		if err != nil {
-			t.Errorf("expected result, got error: %s", err.Error())
-		} else if len(res.Generations) != num {
-			t.Errorf("expected %d gnerations, got %d", num, len(res.Generations))
-		}
-	})
-
-	t.Run("Generate likelihood with generation", func(t *testing.T) {
-		res, err := co.Generate("medium", GenerateOptions{
-			Prompt:            "Hello my name is",
-			MaxTokens:         10,
-			Temperature:       0.75,
-			ReturnLikelihoods: "GENERATION",
-		})
-		if err != nil {
-			t.Errorf("expected result, got error: %s", err.Error())
-		}
-		if res.Generations[0].Likelihood == nil {
-			t.Errorf("expected likelihood")
-		}
-	})
-
-	t.Run("Generate likelihood with all", func(t *testing.T) {
-		res, err := co.Generate("medium", GenerateOptions{
-			Prompt:            "Hello my name is",
-			MaxTokens:         10,
-			Temperature:       0.75,
-			ReturnLikelihoods: "ALL",
-		})
-		if err != nil {
-			t.Errorf("expected result, got error: %s", err.Error())
-		}
-		if res.Generations[0].Likelihood == nil {
-			t.Errorf("expected likelihood")
-		}
-	})
-
-	t.Run("Generate likelihood with none", func(t *testing.T) {
-		res, err := co.Generate("medium", GenerateOptions{
-			Prompt:            "Hello my name is",
-			MaxTokens:         10,
-			Temperature:       0.75,
-			ReturnLikelihoods: "NONE",
-		})
-		if err != nil {
-			t.Errorf("expected result, got error: %s", err.Error())
-		}
-		if res.Generations[0].Likelihood != nil {
-			t.Errorf("expected nil, got %p", res.Generations[0].Likelihood)
-		}
-	})
-}
-
-func TestEmbed(t *testing.T) {
-	co, err := CreateClient(apiKey)
-	if err != nil {
-		t.Error(err)
-	}
-
-	t.Run("Embed", func(t *testing.T) {
-		texts := []string{"hello", "goodbye"}
-
-		_, err := co.Embed("small", EmbedOptions{
-			Texts:    texts,
-			Truncate: TruncateNone,
-		})
-		if err != nil {
-			t.Errorf("expected result, got error: %s", err.Error())
-		}
-	})
-}
-
-func TestTokenize(t *testing.T) {
-	t.Run("TokenizeSuccess", func(t *testing.T) {
-		text := "tokenize me!"
-
-		_, err := Tokenize(TokenizeOptions{
-			Text: text,
-		})
-		if err != nil {
-			t.Errorf("Expected result, got error: %s", err.Error())
-		}
-	})
-
-	t.Run("TokenizeEmptyText", func(t *testing.T) {
-		text := ""
-
-		res, err := Tokenize(TokenizeOptions{
-			Text: text,
-		})
-		if err != nil {
-			t.Errorf("Expected result, got error: %s", err.Error())
-		}
-		expected := []int64{}
-		if len(res.Tokens) != 0 {
-			t.Errorf("Tokenization failed. Expected: %v, Output: %v", res.Tokens, expected)
-		}
-	})
-}
-
-func TestClassify(t *testing.T) {
-	co, err := CreateClient(apiKey)
-	if err != nil {
-		t.Error(err)
-	}
-
-	t.Run("ClassifySuccessMinimumFields", func(t *testing.T) {
-		res, err := co.Classify("medium", ClassifyOptions{
-			Inputs: []string{"purple"},
-			Examples: []Example{
-				{"apple", "fruit"}, {"banana", "fruit"}, {"watermelon", "fruit"}, {"cherry", "fruit"}, {"lemon", "fruit"},
-				{"red", "color"}, {"blue", "color"}, {"blue", "color"}, {"yellow", "color"}, {"green", "color"}},
-		})
-
-		if err != nil {
-			t.Errorf("Expected result, got error: %s", err.Error())
-		}
-
-		if res.Classifications[0].Prediction != "color" {
-			t.Errorf("Expected: color. Receieved: %s", res.Classifications[0].Prediction)
-		}
-	})
-
-	t.Run("ClassifySuccessAllFields", func(t *testing.T) {
-		res, err := co.Classify("medium", ClassifyOptions{
-			TaskDescription: "Classify these words as either a color or a fruit.",
-			Inputs:          []string{"grape", "pink"},
-			Examples: []Example{
-				{"apple", "fruit"}, {"banana", "fruit"}, {"watermelon", "fruit"}, {"cherry", "fruit"}, {"lemon", "fruit"},
-				{"red", "color"}, {"blue", "color"}, {"blue", "color"}, {"yellow", "color"}, {"green", "color"}},
-			OutputIndicator: "This is a",
-		})
-
-		if err != nil {
-			t.Errorf("Expected result, got error: %s", err.Error())
-		}
-
-		if res.Classifications[0].Prediction != "fruit" {
-			t.Errorf("Expected: fruit. Receieved: %s", res.Classifications[0].Prediction)
-		}
-		if res.Classifications[1].Prediction != "color" {
-			t.Errorf("Expected: color. Receieved: %s", res.Classifications[1].Prediction)
-		}
-	})
-
-	t.Run("ClassifySuccessTaskDescription", func(t *testing.T) {
-		res, err := co.Classify("medium", ClassifyOptions{
-			TaskDescription: "Classify these words as a fruit or a color",
-			Inputs:          []string{"kiwi"},
-			Examples: []Example{
-				{"apple", "fruit"}, {"banana", "fruit"}, {"watermelon", "fruit"}, {"cherry", "fruit"}, {"lemon", "fruit"},
-				{"red", "color"}, {"blue", "color"}, {"blue", "color"}, {"yellow", "color"}, {"green", "color"}},
-		})
-
-		if err != nil {
-			t.Errorf("Expected result, got error: %s", err.Error())
-		}
-
-		if res.Classifications[0].Prediction != "fruit" {
-			t.Errorf("Expected: fruit. Receieved: %s", res.Classifications[0].Prediction)
-		}
-	})
-
-	t.Run("ClassifySuccessOutputIndicator", func(t *testing.T) {
-		res, err := co.Classify("medium", ClassifyOptions{
-			Inputs: []string{"pineapple"},
-			Examples: []Example{
-				{"apple", "fruit"}, {"banana", "fruit"}, {"watermelon", "fruit"}, {"cherry", "fruit"}, {"lemon", "fruit"},
-				{"red", "color"}, {"blue", "color"}, {"blue", "color"}, {"yellow", "color"}, {"green", "color"}},
-			OutputIndicator: "This is a",
-		})
-
-		if err != nil {
-			t.Errorf("Expected result, got error: %s", err.Error())
-		}
-
-		if res.Classifications[0].Prediction != "fruit" {
-			t.Errorf("Expected: fruit. Receieved: %s", res.Classifications[0].Prediction)
-		}
-	})
-}
diff --git a/embed_test.go b/embed_test.go
@@ -0,0 +1,24 @@
+package cohere
+
+import (
+	"testing"
+)
+
+func TestEmbed(t *testing.T) {
+	co, err := CreateClient(apiKey)
+	if err != nil {
+		t.Error(err)
+	}
+
+	t.Run("Embed", func(t *testing.T) {
+		texts := []string{"hello", "goodbye"}
+
+		_, err := co.Embed("small", EmbedOptions{
+			Texts:    texts,
+			Truncate: TruncateNone,
+		})
+		if err != nil {
+			t.Errorf("expected result, got error: %s", err.Error())
+		}
+	})
+}