From 020cf3a9ecb8215dc76e47f6f235cd12bdf67c9c Mon Sep 17 00:00:00 2001
From: Roman Glushko <roman.glushko.m@gmail.com>
Date: Mon, 8 Jan 2024 23:26:30 +0200
Subject: [PATCH] #43: Covered router workflow by tests

---
 pkg/routers/router_test.go | 72 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 71 insertions(+), 1 deletion(-)

diff --git a/pkg/routers/router_test.go b/pkg/routers/router_test.go
index 2118eb8a..be053afa 100644
--- a/pkg/routers/router_test.go
+++ b/pkg/routers/router_test.go
@@ -14,7 +14,7 @@ import (
 	"glide/pkg/telemetry"
 )
 
-func TestLangRouter_Priority_ChatRequest(t *testing.T) {
+func TestLangRouter_Priority_PickFistHealthy(t *testing.T) {
 	budget := health.NewErrorBudget(3, health.SEC)
 	models := []*providers.LangModel{
 		providers.NewLangModel(
@@ -49,3 +49,73 @@ func TestLangRouter_Priority_ChatRequest(t *testing.T) {
 		require.NoError(t, err)
 	}
 }
+
+func TestLangRouter_Priority_PickSecondHealthy(t *testing.T) {
+	budget := health.NewErrorBudget(3, health.SEC)
+	models := []*providers.LangModel{
+		providers.NewLangModel(
+			"first",
+			providers.NewProviderMock([]providers.ResponseMock{{Err: &ErrNoModelAvailable}, {Msg: "2"}}),
+			*budget,
+		),
+		providers.NewLangModel(
+			"second",
+			providers.NewProviderMock([]providers.ResponseMock{{Msg: "1"}}),
+			*budget,
+		),
+	}
+
+	expectedModels := []string{"second", "first"}
+
+	router := LangRouter{
+		routerID:  "test_router",
+		Config:    &LangRouterConfig{},
+		retry:     retry.NewExpRetry(3, 2, 1*time.Second, nil),
+		routing:   routing.NewPriorityRouting(models),
+		models:    models,
+		telemetry: telemetry.NewTelemetryMock(),
+	}
+
+	ctx := context.Background()
+	req := schemas.NewChatFromStr("tell me a dad joke")
+
+	for _, modelID := range expectedModels {
+		resp, err := router.Chat(ctx, req)
+
+		require.NoError(t, err)
+		require.Equal(t, modelID, resp.Model)
+		require.Equal(t, "test_router", resp.Router)
+	}
+}
+
+func TestLangRouter_Priority_AllModelsUnavailable(t *testing.T) {
+	budget := health.NewErrorBudget(3, health.SEC)
+	models := []*providers.LangModel{
+		providers.NewLangModel(
+			"first",
+			providers.NewProviderMock([]providers.ResponseMock{{Err: &ErrNoModelAvailable}, {Err: &ErrNoModelAvailable}}),
+			*budget,
+		),
+		providers.NewLangModel(
+			"second",
+			providers.NewProviderMock([]providers.ResponseMock{{Err: &ErrNoModelAvailable}, {Err: &ErrNoModelAvailable}}),
+			*budget,
+		),
+	}
+
+	router := LangRouter{
+		routerID:  "test_router",
+		Config:    &LangRouterConfig{},
+		retry:     retry.NewExpRetry(1, 2, 1*time.Millisecond, nil),
+		routing:   routing.NewPriorityRouting(models),
+		models:    models,
+		telemetry: telemetry.NewTelemetryMock(),
+	}
+
+	ctx := context.Background()
+	req := schemas.NewChatFromStr("tell me a dad joke")
+
+	_, err := router.Chat(ctx, req)
+
+	require.Error(t, err)
+}