From c932f8065977afc4d21793d65c999b754565e963 Mon Sep 17 00:00:00 2001
From: Phil Winder <phil@winder.ai>
Date: Wed, 6 Nov 2024 18:51:14 +0000
Subject: [PATCH] test(openai): fix test

---
 api/pkg/openai/helix_openai_server_test.go | 37 ++++++++++++----------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/api/pkg/openai/helix_openai_server_test.go b/api/pkg/openai/helix_openai_server_test.go
index 397eaf96f..e9066d1d9 100644
--- a/api/pkg/openai/helix_openai_server_test.go
+++ b/api/pkg/openai/helix_openai_server_test.go
@@ -3,6 +3,7 @@ package openai
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/helixml/helix/api/pkg/config"
 	"github.com/helixml/helix/api/pkg/model"
@@ -37,7 +38,7 @@ func (suite *HelixOpenAiServerTestSuite) SetupTest() {
 	suite.pubsub = pubsub
 
 	cfg, _ := config.LoadServerConfig()
-	scheduler := scheduler.NewScheduler(&cfg)
+	scheduler := scheduler.NewScheduler(suite.ctx, &cfg, nil)
 	scheduler.UpdateRunner(&types.RunnerState{
 		ID:          "runner-1",
 		TotalMemory: model.GB * 24, // 24GB runner
@@ -51,22 +52,12 @@ func (suite *HelixOpenAiServerTestSuite) SetupTest() {
 }
 
 func (suite *HelixOpenAiServerTestSuite) Test_GetNextLLMInferenceRequest() {
+	enqueueTestLLMWorkload(suite.srv.scheduler, "req-1", model.Model_Ollama_Llama3_70b)
+	enqueueTestLLMWorkload(suite.srv.scheduler, "req-2", model.Model_Ollama_Llama3_8b)
 
-	// Add a request to the queue
-	suite.srv.queue = append(suite.srv.queue,
-		&types.RunnerLLMInferenceRequest{
-			RequestID: "req-1",
-			Request: &openai.ChatCompletionRequest{
-				Model: model.Model_Ollama_Llama3_70b,
-			},
-		},
-		&types.RunnerLLMInferenceRequest{
-			RequestID: "req-2",
-			Request: &openai.ChatCompletionRequest{
-				Model: model.Model_Ollama_Llama3_8b,
-			},
-		},
-	)
+	// Enough time for the internal goroutine to process the queue. No way of getting access to this
+	// outside of the scheduler.
+	time.Sleep(100 * time.Millisecond)
 
 	req, err := suite.srv.GetNextLLMInferenceRequest(suite.ctx, types.InferenceRequestFilter{}, "runner-1")
 	suite.Require().NoError(err)
@@ -74,3 +65,17 @@ func (suite *HelixOpenAiServerTestSuite) Test_GetNextLLMInferenceRequest() {
 
 	suite.Equal("req-2", req.RequestID)
 }
+
+func enqueueTestLLMWorkload(s scheduler.Scheduler, name string, model string) error {
+	req := &types.RunnerLLMInferenceRequest{
+		RequestID: name,
+		Request: &openai.ChatCompletionRequest{
+			Model: model,
+		},
+	}
+	work, err := scheduler.NewLLMWorkload(req)
+	if err != nil {
+		return err
+	}
+	return s.Enqueue(work)
+}