agents: improve TestConversationalWithMemory testing

raeeceip · Feb 4, 2024 · 8cefb0c · 8cefb0c
1 parent edc9740
commit 8cefb0c
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/agents/conversational_test.go b/agents/conversational_test.go
@@ -3,7 +3,7 @@ package agents
 import (
 	"context"
 	"os"
-	"strings"
+	"regexp"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -35,5 +35,8 @@ func TestConversationalWithMemory(t *testing.T) {
 
 	res, err := chains.Run(context.Background(), executor, "What is the year I was born times 34")
 	require.NoError(t, err)
-	require.True(t, strings.Contains(res, "67558"), `result does not contain the correct answer '67558'`)
+	expectedRe := "67,?558"
+	if !regexp.MustCompile(expectedRe).MatchString(res) {
+		t.Errorf("result does not contain the crrect answer '67558', got: %s", res)
+	}
 }