fix custom dataset test case

Signed-off-by: Qifan Deng <[email protected]>

Author: pancak3

pkg/dataset/custom_dataset.go

@@ -25,7 +25,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"math/rand"
 	"net/http"
 	"os"
 	"os/signal"
@@ -385,6 +384,6 @@ func (d *CustomDataset) GenerateTokens(req openaiserverapi.CompletionRequest, nT
 		return GenPresetRandomTokens(nTokens), nil
 	}
 	d.hasWarned = false
-	randIndex := rand.Intn(len(tokensList))
+	randIndex := common.RandomInt(0, len(tokensList)-1)
 	return tokensList[randIndex], nil
 }

pkg/dataset/custom_dataset_test.go

@@ -19,6 +19,7 @@ package dataset
 import (
 	"encoding/json"
 	"os"
+	"time"
 
 	"github.com/go-logr/logr"
 	"github.com/llm-d/llm-d-inference-sim/pkg/common"
@@ -33,7 +34,7 @@ const (
 	testPrompt = "Hello world!"
 )
 
-var _ = Describe("CustomDataset", func() {
+var _ = Describe("CustomDataset", Ordered, func() {
 	var (
 		dataset               *CustomDataset
 		file_folder           string
@@ -46,6 +47,10 @@ var _ = Describe("CustomDataset", func() {
 		pathToInvalidTypeDB   string
 	)
 
+	BeforeAll(func() {
+		common.InitRandom(time.Now().UnixNano())
+	})
+
 	BeforeEach(func() {
 		dataset = &CustomDataset{
 			BaseDataset: BaseDataset{

pkg/dataset/dataset_test.go

@@ -36,6 +36,7 @@ var _ = Describe("Dataset", Ordered, func() {
 	BeforeAll(func() {
 		common.InitRandom(time.Now().UnixNano())
 	})
+
 	BeforeEach(func() {
 		dataset = &BaseDataset{
 			Logger: logr.Discard(),