Skip to content

Commit

Permalink
Revert "remove test code that's not building"
Browse files Browse the repository at this point in the history
This reverts commit 4bcb63d.
  • Loading branch information
philwinder committed Nov 6, 2024
1 parent 4bcb63d commit b0e4651
Showing 1 changed file with 76 additions and 0 deletions.
76 changes: 76 additions & 0 deletions api/pkg/openai/helix_openai_server_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
package openai

import (
"context"
"testing"

"github.com/helixml/helix/api/pkg/config"
"github.com/helixml/helix/api/pkg/model"
"github.com/helixml/helix/api/pkg/pubsub"
"github.com/helixml/helix/api/pkg/scheduler"
"github.com/helixml/helix/api/pkg/types"
openai "github.com/sashabaranov/go-openai"
"github.com/stretchr/testify/suite"
gomock "go.uber.org/mock/gomock"
)

func TestHelixOpenAiServerTestSuite(t *testing.T) {
suite.Run(t, new(HelixOpenAiServerTestSuite))
}

type HelixOpenAiServerTestSuite struct {
ctx context.Context
suite.Suite
ctrl *gomock.Controller
pubsub pubsub.PubSub

srv *InternalHelixServer
}

func (suite *HelixOpenAiServerTestSuite) SetupTest() {
suite.ctx = context.Background()
suite.ctrl = gomock.NewController(suite.T())

pubsub, err := pubsub.NewInMemoryNats(suite.T().TempDir())
suite.Require().NoError(err)

suite.pubsub = pubsub

cfg, _ := config.LoadServerConfig()
scheduler := scheduler.NewScheduler(&cfg)
scheduler.UpdateRunner(&types.RunnerState{
ID: "runner-1",
TotalMemory: model.GB * 24, // 24GB runner
})
scheduler.UpdateRunner(&types.RunnerState{
ID: "runner-2",
TotalMemory: model.GB * 48, // 48GB runner
})
suite.Require().NoError(err)
suite.srv = NewInternalHelixServer(&cfg, pubsub, scheduler)
}

func (suite *HelixOpenAiServerTestSuite) Test_GetNextLLMInferenceRequest() {

// 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,
},
},
)

req, err := suite.srv.GetNextLLMInferenceRequest(suite.ctx, types.InferenceRequestFilter{}, "runner-1")
suite.Require().NoError(err)
suite.Require().NotNil(req)

suite.Equal("req-2", req.RequestID)
}

0 comments on commit b0e4651

Please sign in to comment.