Nik Afiq
ad50d641bd
- Updated LLMClient interface to support model-specific generation and model listing. - Integrated model store and validator into the command application for managing AI models. - Implemented commands for setting, getting, and listing active AI models in Discord. - Enhanced AI query handling to utilize the selected model and return model information in responses. - Added caching mechanism for model validation to improve performance. - Introduced gRPC methods for listing available AI models in the ai-gateway. - Updated protobuf definitions to include model-related fields and messages. - Added tests for model store and validator functionalities.
550 lines
15 KiB
Go
550 lines
15 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"reflect"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.nik4nao.com/nik/home-services/discord-bot/internal/core/ports/driven"
|
|
"gitea.nik4nao.com/nik/home-services/discord-bot/internal/modelstore"
|
|
"gitea.nik4nao.com/nik/home-services/discord-bot/internal/modelvalidator"
|
|
)
|
|
|
|
type mockHAGateway struct {
|
|
listLightsFunc func(ctx context.Context) ([]driven.Light, error)
|
|
listSwitchesFunc func(ctx context.Context) ([]driven.Switch, error)
|
|
turnOnLightFunc func(ctx context.Context, entityID string, brightnessPct *uint32, colorTempKelvin *uint32) error
|
|
turnOffLightFunc func(ctx context.Context, entityID string, transition *uint32) error
|
|
toggleLightFunc func(ctx context.Context, entityID string) error
|
|
}
|
|
|
|
type mockAIGateway struct {
|
|
queryFunc func(ctx context.Context, text, model string) (string, string, error)
|
|
listModelsFunc func(ctx context.Context) ([]string, error)
|
|
}
|
|
|
|
func (m *mockHAGateway) ListLights(ctx context.Context) ([]driven.Light, error) {
|
|
if m.listLightsFunc == nil {
|
|
return nil, nil
|
|
}
|
|
return m.listLightsFunc(ctx)
|
|
}
|
|
|
|
func (m *mockHAGateway) ListSwitches(ctx context.Context) ([]driven.Switch, error) {
|
|
if m.listSwitchesFunc == nil {
|
|
return nil, nil
|
|
}
|
|
return m.listSwitchesFunc(ctx)
|
|
}
|
|
|
|
func (m *mockHAGateway) TurnOnLight(ctx context.Context, entityID string, brightnessPct *uint32, colorTempKelvin *uint32) error {
|
|
if m.turnOnLightFunc == nil {
|
|
return nil
|
|
}
|
|
return m.turnOnLightFunc(ctx, entityID, brightnessPct, colorTempKelvin)
|
|
}
|
|
|
|
func (m *mockHAGateway) TurnOffLight(ctx context.Context, entityID string, transition *uint32) error {
|
|
if m.turnOffLightFunc == nil {
|
|
return nil
|
|
}
|
|
return m.turnOffLightFunc(ctx, entityID, transition)
|
|
}
|
|
|
|
func (m *mockHAGateway) ToggleLight(ctx context.Context, entityID string) error {
|
|
if m.toggleLightFunc == nil {
|
|
return nil
|
|
}
|
|
return m.toggleLightFunc(ctx, entityID)
|
|
}
|
|
|
|
func (m *mockAIGateway) Query(ctx context.Context, text, model string) (string, string, error) {
|
|
if m.queryFunc == nil {
|
|
return "", "", nil
|
|
}
|
|
return m.queryFunc(ctx, text, model)
|
|
}
|
|
|
|
func (m *mockAIGateway) ListModels(ctx context.Context) ([]string, error) {
|
|
if m.listModelsFunc == nil {
|
|
return nil, nil
|
|
}
|
|
return m.listModelsFunc(ctx)
|
|
}
|
|
|
|
func newTestCommandApp(ha *mockHAGateway, ai *mockAIGateway) *CommandApp {
|
|
return NewCommandApp(ha, ai, modelstore.New(), modelvalidator.New(ai, time.Minute))
|
|
}
|
|
|
|
func TestCommandAppHandleLightList(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
lights []driven.Light
|
|
want string
|
|
}{
|
|
{
|
|
name: "empty list",
|
|
want: "No lights found.",
|
|
},
|
|
{
|
|
name: "single light with full data",
|
|
lights: []driven.Light{
|
|
{
|
|
EntityID: "light.kitchen",
|
|
FriendlyName: "Kitchen",
|
|
State: "on",
|
|
SupportedColorModes: []string{"brightness", "color_temp"},
|
|
MinColorTempKelvin: 2700,
|
|
MaxColorTempKelvin: 6500,
|
|
IsHueGroup: true,
|
|
},
|
|
},
|
|
want: "```text\n🟢 Kitchen on brightness,color_temp 2700-6500K hue-group\n```",
|
|
},
|
|
{
|
|
name: "empty friendly name uses entity id and off emoji",
|
|
lights: []driven.Light{
|
|
{
|
|
EntityID: "light.desk",
|
|
State: "off",
|
|
},
|
|
},
|
|
want: "```text\n🔴 light.desk off -\n```",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
app := newTestCommandApp(&mockHAGateway{
|
|
listLightsFunc: func(ctx context.Context) ([]driven.Light, error) {
|
|
return tt.lights, nil
|
|
},
|
|
}, &mockAIGateway{})
|
|
|
|
got, err := app.HandleLightList(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("HandleLightList() error = %v", err)
|
|
}
|
|
if got != tt.want {
|
|
t.Fatalf("HandleLightList() = %q, want %q", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCommandAppHandleLightOn(t *testing.T) {
|
|
brightness := uint32(55)
|
|
colorTemp := uint32(3200)
|
|
|
|
tests := []struct {
|
|
name string
|
|
lights []driven.Light
|
|
entityID string
|
|
brightnessPct *uint32
|
|
colorTempKelvin *uint32
|
|
turnOnErr error
|
|
listErr error
|
|
want string
|
|
wantErr string
|
|
}{
|
|
{
|
|
name: "no optional params",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
entityID: "light.kitchen",
|
|
want: "Turned on `Kitchen`.",
|
|
},
|
|
{
|
|
name: "brightness only",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
entityID: "light.kitchen",
|
|
brightnessPct: &brightness,
|
|
want: "Turned on `Kitchen` (brightness 55%).",
|
|
},
|
|
{
|
|
name: "color temp only",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
entityID: "light.kitchen",
|
|
colorTempKelvin: &colorTemp,
|
|
want: "Turned on `Kitchen` (3200K).",
|
|
},
|
|
{
|
|
name: "both params",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
entityID: "light.kitchen",
|
|
brightnessPct: &brightness,
|
|
colorTempKelvin: &colorTemp,
|
|
want: "Turned on `Kitchen` (brightness 55%, 3200K).",
|
|
},
|
|
{
|
|
name: "light not found falls back to entity id",
|
|
lights: []driven.Light{{EntityID: "light.other", FriendlyName: "Other"}},
|
|
entityID: "light.kitchen",
|
|
want: "Turned on `light.kitchen`.",
|
|
},
|
|
{
|
|
name: "TurnOnLight error",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
entityID: "light.kitchen",
|
|
turnOnErr: errors.New("boom"),
|
|
wantErr: "handle light on: boom",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
var gotEntityID string
|
|
var gotBrightness *uint32
|
|
var gotColorTemp *uint32
|
|
|
|
app := newTestCommandApp(&mockHAGateway{
|
|
listLightsFunc: func(ctx context.Context) ([]driven.Light, error) {
|
|
if tt.listErr != nil {
|
|
return nil, tt.listErr
|
|
}
|
|
return tt.lights, nil
|
|
},
|
|
turnOnLightFunc: func(ctx context.Context, entityID string, brightnessPct *uint32, colorTempKelvin *uint32) error {
|
|
gotEntityID = entityID
|
|
gotBrightness = brightnessPct
|
|
gotColorTemp = colorTempKelvin
|
|
return tt.turnOnErr
|
|
},
|
|
}, &mockAIGateway{})
|
|
|
|
got, err := app.HandleLightOn(context.Background(), tt.entityID, tt.brightnessPct, tt.colorTempKelvin)
|
|
if tt.wantErr != "" {
|
|
if err == nil || err.Error() != tt.wantErr {
|
|
t.Fatalf("HandleLightOn() error = %v, want %q", err, tt.wantErr)
|
|
}
|
|
return
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("HandleLightOn() error = %v", err)
|
|
}
|
|
if got != tt.want {
|
|
t.Fatalf("HandleLightOn() = %q, want %q", got, tt.want)
|
|
}
|
|
if gotEntityID != tt.entityID {
|
|
t.Fatalf("TurnOnLight entityID = %q, want %q", gotEntityID, tt.entityID)
|
|
}
|
|
if !reflect.DeepEqual(gotBrightness, tt.brightnessPct) || !reflect.DeepEqual(gotColorTemp, tt.colorTempKelvin) {
|
|
t.Fatalf("TurnOnLight params = (%v, %v), want (%v, %v)", gotBrightness, gotColorTemp, tt.brightnessPct, tt.colorTempKelvin)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCommandAppHandleLightOff(t *testing.T) {
|
|
transition := uint32(4)
|
|
|
|
tests := []struct {
|
|
name string
|
|
entityID string
|
|
lights []driven.Light
|
|
transition *uint32
|
|
turnOffErr error
|
|
want string
|
|
wantErr string
|
|
}{
|
|
{
|
|
name: "no transition",
|
|
entityID: "light.kitchen",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
want: "Turned off `Kitchen`.",
|
|
},
|
|
{
|
|
name: "with transition",
|
|
entityID: "light.kitchen",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
transition: &transition,
|
|
want: "Turned off `Kitchen` with 4s transition.",
|
|
},
|
|
{
|
|
name: "error",
|
|
entityID: "light.kitchen",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
turnOffErr: errors.New("boom"),
|
|
wantErr: "handle light off: boom",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
var gotTransition *uint32
|
|
app := newTestCommandApp(&mockHAGateway{
|
|
listLightsFunc: func(ctx context.Context) ([]driven.Light, error) {
|
|
return tt.lights, nil
|
|
},
|
|
turnOffLightFunc: func(ctx context.Context, entityID string, transition *uint32) error {
|
|
gotTransition = transition
|
|
return tt.turnOffErr
|
|
},
|
|
}, &mockAIGateway{})
|
|
|
|
got, err := app.HandleLightOff(context.Background(), tt.entityID, tt.transition)
|
|
if tt.wantErr != "" {
|
|
if err == nil || err.Error() != tt.wantErr {
|
|
t.Fatalf("HandleLightOff() error = %v, want %q", err, tt.wantErr)
|
|
}
|
|
return
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("HandleLightOff() error = %v", err)
|
|
}
|
|
if got != tt.want {
|
|
t.Fatalf("HandleLightOff() = %q, want %q", got, tt.want)
|
|
}
|
|
if !reflect.DeepEqual(gotTransition, tt.transition) {
|
|
t.Fatalf("TurnOffLight transition = %v, want %v", gotTransition, tt.transition)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCommandAppHandleLightToggle(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
lights []driven.Light
|
|
listErr error
|
|
toggleErr error
|
|
want string
|
|
wantErr string
|
|
}{
|
|
{
|
|
name: "happy path",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
want: "Toggled `Kitchen`.",
|
|
},
|
|
{
|
|
name: "lookup error",
|
|
listErr: errors.New("lookup failed"),
|
|
wantErr: "lookup light name: list lights: lookup failed",
|
|
},
|
|
{
|
|
name: "toggle error",
|
|
lights: []driven.Light{{EntityID: "light.kitchen", FriendlyName: "Kitchen"}},
|
|
toggleErr: errors.New("boom"),
|
|
wantErr: "handle light toggle: boom",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
app := newTestCommandApp(&mockHAGateway{
|
|
listLightsFunc: func(ctx context.Context) ([]driven.Light, error) {
|
|
if tt.listErr != nil {
|
|
return nil, tt.listErr
|
|
}
|
|
return tt.lights, nil
|
|
},
|
|
toggleLightFunc: func(ctx context.Context, entityID string) error {
|
|
return tt.toggleErr
|
|
},
|
|
}, &mockAIGateway{})
|
|
|
|
got, err := app.HandleLightToggle(context.Background(), "light.kitchen")
|
|
if tt.wantErr != "" {
|
|
if err == nil || err.Error() != tt.wantErr {
|
|
t.Fatalf("HandleLightToggle() error = %v, want %q", err, tt.wantErr)
|
|
}
|
|
return
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("HandleLightToggle() error = %v", err)
|
|
}
|
|
if got != tt.want {
|
|
t.Fatalf("HandleLightToggle() = %q, want %q", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCommandAppHandleSwitchList(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
switches []driven.Switch
|
|
want string
|
|
}{
|
|
{
|
|
name: "empty",
|
|
want: "No switches found.",
|
|
},
|
|
{
|
|
name: "with switches",
|
|
switches: []driven.Switch{
|
|
{EntityID: "switch.fan", FriendlyName: "Fan", State: "on", DeviceClass: "outlet"},
|
|
{EntityID: "switch.pump", State: "off"},
|
|
},
|
|
want: "```text\n🟢 Fan on outlet\n🔴 switch.pump off switch\n```",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
app := newTestCommandApp(&mockHAGateway{
|
|
listSwitchesFunc: func(ctx context.Context) ([]driven.Switch, error) {
|
|
return tt.switches, nil
|
|
},
|
|
}, &mockAIGateway{})
|
|
|
|
got, err := app.HandleSwitchList(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("HandleSwitchList() error = %v", err)
|
|
}
|
|
if got != tt.want {
|
|
t.Fatalf("HandleSwitchList() = %q, want %q", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCommandAppAutocompleteLights(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
lights []driven.Light
|
|
listErr error
|
|
want []Choice
|
|
wantErr string
|
|
}{
|
|
{
|
|
name: "friendly name and fallback",
|
|
lights: []driven.Light{
|
|
{EntityID: "light.kitchen", FriendlyName: "Kitchen"},
|
|
{EntityID: "light.desk"},
|
|
},
|
|
want: []Choice{
|
|
{Label: "Kitchen", Value: "light.kitchen"},
|
|
{Label: "light.desk", Value: "light.desk"},
|
|
},
|
|
},
|
|
{
|
|
name: "ListLights error",
|
|
listErr: errors.New("boom"),
|
|
wantErr: "autocomplete lights: boom",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
app := newTestCommandApp(&mockHAGateway{
|
|
listLightsFunc: func(ctx context.Context) ([]driven.Light, error) {
|
|
if tt.listErr != nil {
|
|
return nil, tt.listErr
|
|
}
|
|
return tt.lights, nil
|
|
},
|
|
}, &mockAIGateway{})
|
|
|
|
got, err := app.AutocompleteLights(context.Background())
|
|
if tt.wantErr != "" {
|
|
if err == nil || err.Error() != tt.wantErr {
|
|
t.Fatalf("AutocompleteLights() error = %v, want %q", err, tt.wantErr)
|
|
}
|
|
return
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("AutocompleteLights() error = %v", err)
|
|
}
|
|
if !reflect.DeepEqual(got, tt.want) {
|
|
t.Fatalf("AutocompleteLights() = %#v, want %#v", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCommandAppAutocompleteSwitches(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
switches []driven.Switch
|
|
listErr error
|
|
want []Choice
|
|
wantErr string
|
|
}{
|
|
{
|
|
name: "friendly name and fallback",
|
|
switches: []driven.Switch{
|
|
{EntityID: "switch.fan", FriendlyName: "Fan"},
|
|
{EntityID: "switch.pump"},
|
|
},
|
|
want: []Choice{
|
|
{Label: "Fan", Value: "switch.fan"},
|
|
{Label: "switch.pump", Value: "switch.pump"},
|
|
},
|
|
},
|
|
{
|
|
name: "ListSwitches error",
|
|
listErr: errors.New("boom"),
|
|
wantErr: "autocomplete switches: boom",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
app := newTestCommandApp(&mockHAGateway{
|
|
listSwitchesFunc: func(ctx context.Context) ([]driven.Switch, error) {
|
|
if tt.listErr != nil {
|
|
return nil, tt.listErr
|
|
}
|
|
return tt.switches, nil
|
|
},
|
|
}, &mockAIGateway{})
|
|
|
|
got, err := app.AutocompleteSwitches(context.Background())
|
|
if tt.wantErr != "" {
|
|
if err == nil || err.Error() != tt.wantErr {
|
|
t.Fatalf("AutocompleteSwitches() error = %v, want %q", err, tt.wantErr)
|
|
}
|
|
return
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("AutocompleteSwitches() error = %v", err)
|
|
}
|
|
if !reflect.DeepEqual(got, tt.want) {
|
|
t.Fatalf("AutocompleteSwitches() = %#v, want %#v", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCommandAppHandleAIQuery(t *testing.T) {
|
|
store := modelstore.New()
|
|
store.Set("llama3:latest")
|
|
app := NewCommandApp(&mockHAGateway{}, &mockAIGateway{
|
|
queryFunc: func(ctx context.Context, text, model string) (string, string, error) {
|
|
if text != "turn on kitchen" {
|
|
t.Fatalf("Query() text = %q", text)
|
|
}
|
|
if model != "llama3:latest" {
|
|
t.Fatalf("Query() model = %q", model)
|
|
}
|
|
return "Turning on Kitchen.", "llama3:latest", nil
|
|
},
|
|
}, store, modelvalidator.New(&mockAIGateway{}, time.Minute))
|
|
|
|
got, err := app.HandleAIQuery(context.Background(), "turn on kitchen")
|
|
if err != nil {
|
|
t.Fatalf("HandleAIQuery() error = %v", err)
|
|
}
|
|
if got != "Turning on Kitchen.\n\n_(via llama3:latest)_" {
|
|
t.Fatalf("HandleAIQuery() = %q", got)
|
|
}
|
|
}
|
|
|
|
func TestCommandAppHandleAIModelSet(t *testing.T) {
|
|
app := newTestCommandApp(&mockHAGateway{}, &mockAIGateway{
|
|
listModelsFunc: func(ctx context.Context) ([]string, error) {
|
|
return []string{"llama3:latest"}, nil
|
|
},
|
|
})
|
|
|
|
got, err := app.HandleAIModelSet(context.Background(), "llama3")
|
|
if err != nil {
|
|
t.Fatalf("HandleAIModelSet() error = %v", err)
|
|
}
|
|
if got != "Active model set to `llama3:latest`." {
|
|
t.Fatalf("HandleAIModelSet() = %q", got)
|
|
}
|
|
}
|