Fixed dataset test (#246)

Signed-off-by: irar2 <[email protected]>

Author: irar2

pkg/dataset/custom_dataset_test.go

@@ -186,7 +186,11 @@ var _ = Describe("CustomDataset", Ordered, func() {
 		tokens, finishReason, err := dataset.GetTokens(req, common.ModeRandom, random)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(finishReason).To(Equal(StopFinishReason))
-		Expect(tokens).To(Equal([]string{"Hello", " llm-d ", "world", "!"}))
+		if len(tokens) >= 4 {
+			// The number of tokens to generate is random, and if it's less than 4
+			// we will not get these tokens
+			Expect(tokens).To(Equal([]string{"Hello", " llm-d ", "world", "!"}))
+		}
 	})
 
 	It("should return at most 2 tokens for existing prompt", func() {