diff --git a/pkg/display/display.go b/pkg/display/display.go index 6dcc5a07..e9dfaa1f 100644 --- a/pkg/display/display.go +++ b/pkg/display/display.go @@ -669,11 +669,25 @@ func (s *Service) handleWSMessage(msg WSMessage) core.Result { )) case "contextmenu:add": name, _ := msg.Data["name"].(string) - marshalResult := core.JSONMarshal(msg.Data["menu"]) + menuValue, ok := msg.Data["menu"] + if !ok || menuValue == nil { + return core.Result{Value: coreerr.E("display.handleWSMessage", "missing required field \"menu\"", nil), OK: false} + } + marshalResult := core.JSONMarshal(menuValue) + if !marshalResult.OK { + if err, ok := marshalResult.Value.(error); ok { + return core.Result{Value: coreerr.E("display.handleWSMessage", "failed to marshal menu definition", err), OK: false} + } + return core.Result{Value: coreerr.E("display.handleWSMessage", "failed to marshal menu definition", nil), OK: false} + } var menuDef contextmenu.ContextMenuDef - if marshalResult.OK { - menuJSON, _ := marshalResult.Value.([]byte) - core.JSONUnmarshal(menuJSON, &menuDef) + menuJSON, _ := marshalResult.Value.([]byte) + unmarshalResult := core.JSONUnmarshal(menuJSON, &menuDef) + if !unmarshalResult.OK { + if err, ok := unmarshalResult.Value.(error); ok { + return core.Result{Value: coreerr.E("display.handleWSMessage", "failed to unmarshal menu definition", err), OK: false} + } + return core.Result{Value: coreerr.E("display.handleWSMessage", "failed to unmarshal menu definition", nil), OK: false} } return c.Action("contextmenu.add").Run(ctx, core.NewOptions( core.Option{Key: "task", Value: contextmenu.TaskAdd{Name: name, Menu: menuDef}}, diff --git a/pkg/display/display_contextmenu_test.go b/pkg/display/display_contextmenu_test.go new file mode 100644 index 00000000..345c2848 --- /dev/null +++ b/pkg/display/display_contextmenu_test.go @@ -0,0 +1,85 @@ +package display + +import ( + "context" + "sync" + "testing" + + core "dappco.re/go/core" + "forge.lthn.ai/core/gui/pkg/contextmenu" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type wsContextMenuPlatform struct { + mu sync.Mutex + menus map[string]contextmenu.ContextMenuDef +} + +func newWSContextMenuPlatform() *wsContextMenuPlatform { + return &wsContextMenuPlatform{ + menus: make(map[string]contextmenu.ContextMenuDef), + } +} + +func (m *wsContextMenuPlatform) Add(name string, menu contextmenu.ContextMenuDef, _ func(string, string, string)) error { + m.mu.Lock() + defer m.mu.Unlock() + m.menus[name] = menu + return nil +} + +func (m *wsContextMenuPlatform) Remove(name string) error { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.menus, name) + return nil +} + +func (m *wsContextMenuPlatform) Get(name string) (*contextmenu.ContextMenuDef, bool) { + m.mu.Lock() + defer m.mu.Unlock() + menu, ok := m.menus[name] + if !ok { + return nil, false + } + return &menu, true +} + +func (m *wsContextMenuPlatform) GetAll() map[string]contextmenu.ContextMenuDef { + m.mu.Lock() + defer m.mu.Unlock() + out := make(map[string]contextmenu.ContextMenuDef, len(m.menus)) + for name, menu := range m.menus { + out[name] = menu + } + return out +} + +func newDisplayWithContextMenu(t *testing.T, platform *wsContextMenuPlatform) (*Service, *core.Core) { + t.Helper() + c := core.New( + core.WithService(Register(nil)), + core.WithService(contextmenu.Register(platform)), + core.WithServiceLock(), + ) + require.True(t, c.ServiceStartup(context.Background(), nil).OK) + return core.MustServiceFor[*Service](c, "display"), c +} + +func TestDisplay_handleWSMessage_ContextMenuAdd_MissingMenu(t *testing.T) { + platform := newWSContextMenuPlatform() + svc, _ := newDisplayWithContextMenu(t, platform) + + result := svc.handleWSMessage(WSMessage{ + Action: "contextmenu:add", + Data: map[string]any{ + "name": "menu", + }, + }) + + require.False(t, result.OK) + assert.Contains(t, result.Value.(error).Error(), `missing required field "menu"`) + _, ok := platform.Get("menu") + assert.False(t, ok) +}