#163: Fixed tests

EinStack · Mar 11, 2024 · d48724d · d48724d
1 parent ba56d3f
commit d48724d
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 11 deletions.
diff --git a/pkg/providers/lang.go b/pkg/providers/lang.go
@@ -2,9 +2,10 @@ package providers
 
 import (
 	"context"
-	"glide/pkg/routers/health"
 	"time"
 
+	"glide/pkg/routers/health"
+
 	"glide/pkg/api/schemas"
 	"glide/pkg/providers/clients"
 	"glide/pkg/routers/latency"
@@ -33,7 +34,7 @@ type LanguageModel struct {
 	modelID               string
 	weight                int
 	client                LangProvider
-	healthTracker         *health.HealthTracker
+	healthTracker         *health.Tracker
 	chatLatency           *latency.MovingAverage
 	chatStreamLatency     *latency.MovingAverage
 	latencyUpdateInterval *time.Duration
@@ -43,7 +44,7 @@ func NewLangModel(modelID string, client LangProvider, budget health.ErrorBudget
 	return &LanguageModel{
 		modelID:               modelID,
 		client:                client,
-		healthTracker:         health.NewHealthTracker(budget),
+		healthTracker:         health.NewTracker(budget),
 		chatLatency:           latency.NewMovingAverage(latencyConfig.Decay, latencyConfig.WarmupSamples),
 		chatStreamLatency:     latency.NewMovingAverage(latencyConfig.Decay, latencyConfig.WarmupSamples),
 		latencyUpdateInterval: latencyConfig.UpdateInterval,

diff --git a/pkg/providers/testing.go b/pkg/providers/testing.go
@@ -104,14 +104,14 @@ type LangModelMock struct {
 	weight      int
 }
 
-func NewLangModelMock(ID string, healthy bool, avgLatency float64, weight int) *LangModelMock {
+func NewLangModelMock(ID string, healthy bool, avgLatency float64, weight int) LangModelMock {
 	chatLatency := latency.NewMovingAverage(0.06, 3)
 
 	if avgLatency > 0.0 {
 		chatLatency.Set(avgLatency)
 	}
 
-	return &LangModelMock{
+	return LangModelMock{
 		modelID:     ID,
 		healthy:     healthy,
 		chatLatency: chatLatency,

diff --git a/pkg/routers/health/tracker.go b/pkg/routers/health/tracker.go
@@ -2,26 +2,28 @@ package health
 
 import (
 	"errors"
+
 	"glide/pkg/providers/clients"
 )
 
-type HealthTracker struct {
+// Tracker tracks errors and general health of model provider
+type Tracker struct {
 	errBudget *TokenBucket
 	rateLimit *RateLimitTracker
 }
 
-func NewHealthTracker(budget ErrorBudget) *HealthTracker {
-	return &HealthTracker{
+func NewTracker(budget ErrorBudget) *Tracker {
+	return &Tracker{
 		rateLimit: NewRateLimitTracker(),
 		errBudget: NewTokenBucket(budget.TimePerTokenMicro(), budget.Budget()),
 	}
 }
 
-func (t *HealthTracker) Healthy() bool {
+func (t *Tracker) Healthy() bool {
 	return !t.rateLimit.Limited() && t.errBudget.HasTokens()
 }
 
-func (t *HealthTracker) TrackErr(err error) {
+func (t *Tracker) TrackErr(err error) {
 	var rateLimitErr *clients.RateLimitError
 
 	if errors.As(err, &rateLimitErr) {

diff --git a/pkg/routers/routing/least_latency_test.go b/pkg/routers/routing/least_latency_test.go
@@ -115,7 +115,8 @@ func TestLeastLatencyRouting_Routing(t *testing.T) {
 			}
 
 			routing := LeastLatencyRouting{
-				schedules: schedules,
+				latencyGetter: providers.ChatMockLatency,
+				schedules:     schedules,
 			}
 
 			iterator := routing.Iterator()