Function TestGetAnswer

modules/generative-cohere/clients/cohere_test.go:37–132

(t *testing.T)

}

func TestGetAnswer(t *testing.T) {
	props := []*modulecapabilities.GenerateProperties{{Text: map[string]string{"prop": "My name is john"}}}

	tests := []struct {
		name           string
		answer         generateResponse
		timeout        time.Duration
		expectedResult string
	}{
		{
			name: "when the server has a successful aner",
			answer: generateResponse{
				Message: responseMessage{
					Content: []responseContent{
						{Text: "John"},
					},
				},
			},
			expectedResult: "John",
		},
		{
			name: "when the server has a an error",
			answer: generateResponse{
				Message: responseMessage{
					ErrorMessage: "some error from the server",
				},
			},
		},
		{
			name: "when the server does not respond in time",
			answer: generateResponse{
				Message: responseMessage{
					ErrorMessage: "context deadline exceeded",
				},
			},
			timeout: time.Second,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			handler := &testAnswerHandler{
				t:       t,
				answer:  test.answer,
				timeout: test.timeout,
			}
			server := httptest.NewServer(handler)
			defer server.Close()

			c := New("apiKey", test.timeout, nullLogger())

			cfg := &fakeClassConfig{baseURL: server.URL}
			res, err := c.GenerateAllResults(context.Background(), props, "What is my name?", nil, false, cfg)

			if test.answer.Message.ErrorMessage != "" {
				assert.Contains(t, err.Error(), test.answer.Message.GetMessage())
			} else {
				assert.Equal(t, test.expectedResult, *res.Result)
			}