robinhood-agentic-mcp/watchlists/watchlists_test.go

package watchlists_test

import (
	"context"
	"encoding/json"
	"sort"
	"testing"

	"github.com/google/go-cmp/cmp"
	"s1d3sw1ped/robinhood-agentic-mcp/client"
	"s1d3sw1ped/robinhood-agentic-mcp/internal/rhntest"
	"s1d3sw1ped/robinhood-agentic-mcp/watchlists"
)

func TestWatchlists_toolNames(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name     string
		call     func(*watchlists.Client) error
		wantName string
		wantArgs map[string]any
	}{
		{
			name: "Lists",
			call: func(c *watchlists.Client) error {
				_, err := c.Lists(context.Background(), watchlists.ListsRequest{})
				return err
			},
			wantName: "get_watchlists",
			wantArgs: map[string]any{},
		},
		{
			name: "Items",
			call: func(c *watchlists.Client) error {
				_, err := c.Items(context.Background(), watchlists.ItemsRequest{ListID: "wl-1"})
				return err
			},
			wantName: "get_watchlist_items",
			wantArgs: map[string]any{"list_id": "wl-1"},
		},
		{
			name: "OptionList",
			call: func(c *watchlists.Client) error {
				_, err := c.OptionList(context.Background(), watchlists.OptionListRequest{})
				return err
			},
			wantName: "get_option_watchlist",
			wantArgs: map[string]any{},
		},
		{
			name: "Popular",
			call: func(c *watchlists.Client) error {
				_, err := c.Popular(context.Background(), watchlists.PopularRequest{})
				return err
			},
			wantName: "get_popular_watchlists",
			wantArgs: map[string]any{},
		},
		{
			name: "Create",
			call: func(c *watchlists.Client) error {
				return c.Create(context.Background(), watchlists.CreateRequest{
					DisplayName:        "Tech",
					IconEmoji:          "📈",
					DisplayDescription: "vwap book",
				})
			},
			wantName: "create_watchlist",
			wantArgs: map[string]any{
				"display_name":        "Tech",
				"icon_emoji":          "📈",
				"display_description": "vwap book",
			},
		},
		{
			name: "Update",
			call: func(c *watchlists.Client) error {
				return c.Update(context.Background(), watchlists.UpdateRequest{
					ListID:             "wl-1",
					DisplayName:        "Tech",
					IconEmoji:          "📈",
					DisplayDescription: "vwap book",
				})
			},
			wantName: "update_watchlist",
			wantArgs: map[string]any{
				"list_id":             "wl-1",
				"display_name":        "Tech",
				"icon_emoji":          "📈",
				"display_description": "vwap book",
			},
		},
		{
			name: "Follow",
			call: func(c *watchlists.Client) error {
				return c.Follow(context.Background(), watchlists.ListIDRequest{ListID: "wl-pop"})
			},
			wantName: "follow_watchlist",
			wantArgs: map[string]any{"list_id": "wl-pop"},
		},
		{
			name: "Unfollow",
			call: func(c *watchlists.Client) error {
				return c.Unfollow(context.Background(), watchlists.ListIDRequest{ListID: "wl-pop"})
			},
			wantName: "unfollow_watchlist",
			wantArgs: map[string]any{"list_id": "wl-pop"},
		},
		{
			name: "Add",
			call: func(c *watchlists.Client) error {
				return c.Add(context.Background(), watchlists.AddRequest{
					ListID:          "wl-1",
					Symbols:         []string{"MU", "SPY"},
					CurrencyPairIDs: []string{"btc-1"},
					IndexIDs:        []string{"idx-spx"},
				})
			},
			wantName: "add_to_watchlist",
			wantArgs: map[string]any{
				"list_id":           "wl-1",
				"symbols":           []string{"MU", "SPY"},
				"currency_pair_ids": []string{"btc-1"},
				"index_ids":         []string{"idx-spx"},
			},
		},
		{
			name: "Remove",
			call: func(c *watchlists.Client) error {
				return c.Remove(context.Background(), watchlists.RemoveRequest{
					ListID:          "wl-1",
					Symbols:         []string{"MU"},
					CurrencyPairIDs: []string{"btc-1"},
					IndexIDs:        []string{"idx-spx"},
				})
			},
			wantName: "remove_from_watchlist",
			wantArgs: map[string]any{
				"list_id":           "wl-1",
				"symbols":           []string{"MU"},
				"currency_pair_ids": []string{"btc-1"},
				"index_ids":         []string{"idx-spx"},
			},
		},
		{
			name: "AddOption",
			call: func(c *watchlists.Client) error {
				return c.AddOption(context.Background(), watchlists.OptionMutateRequest{
					OptionIDs:    []string{"opt-1", "opt-2"},
					PositionType: "long",
				})
			},
			wantName: "add_option_to_watchlist",
			wantArgs: map[string]any{
				"option_ids":    []string{"opt-1", "opt-2"},
				"position_type": "long",
			},
		},
		{
			name: "RemoveOption",
			call: func(c *watchlists.Client) error {
				return c.RemoveOption(context.Background(), watchlists.OptionMutateRequest{
					OptionIDs:    []string{"opt-1"},
					PositionType: "short",
				})
			},
			wantName: "remove_option_from_watchlist",
			wantArgs: map[string]any{
				"option_ids":    []string{"opt-1"},
				"position_type": "short",
			},
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			var gotName string
			var gotArgs map[string]any
			c := watchlists.New(client.Func(func(ctx context.Context, name string, args map[string]any) (json.RawMessage, error) {
				gotName, gotArgs = name, args
				return json.RawMessage(`{}`), nil
			}))
			if err := tc.call(c); err != nil {
				t.Fatal(err)
			}
			if gotName != tc.wantName {
				t.Fatalf("%s %+v", gotName, gotArgs)
			}
			if diff := cmp.Diff(tc.wantArgs, gotArgs); diff != "" {
				t.Fatal(diff)
			}
		})
	}
}

func TestTools(t *testing.T) {
	t.Parallel()
	want := []string{
		"add_option_to_watchlist",
		"add_to_watchlist",
		"create_watchlist",
		"follow_watchlist",
		"get_option_watchlist",
		"get_popular_watchlists",
		"get_watchlist_items",
		"get_watchlists",
		"remove_from_watchlist",
		"remove_option_from_watchlist",
		"unfollow_watchlist",
		"update_watchlist",
	}
	got := append([]string(nil), watchlists.Tools()...)
	sort.Strings(got)
	if diff := cmp.Diff(want, got); diff != "" {
		t.Fatal(diff)
	}
}

func TestLists_rhntest(t *testing.T) {
	t.Parallel()
	s := rhntest.New(t)
	s.Set("get_watchlists", json.RawMessage(`{"watchlists":[{"id":"wl-1","title":"Tech"}]}`))
	c := watchlists.New(&client.Client{URL: s.URL})
	got, err := c.Lists(context.Background(), watchlists.ListsRequest{})
	if err != nil {
		t.Fatal(err)
	}
	if len(got.Watchlists) != 1 || got.Watchlists[0].ID != "wl-1" || got.Watchlists[0].Title != "Tech" {
		t.Fatalf("%+v", got)
	}
	if s.LastName() != "get_watchlists" {
		t.Fatalf("%s", s.LastName())
	}
	if diff := cmp.Diff(map[string]any{}, s.LastArgs()); diff != "" {
		t.Fatal(diff)
	}
}

func TestItems_rhntest(t *testing.T) {
	t.Parallel()
	s := rhntest.New(t)
	s.Set("get_watchlist_items", json.RawMessage(`{"items":[{"symbol":"MU","object_type":"equity"}]}`))
	c := watchlists.New(&client.Client{URL: s.URL})
	got, err := c.Items(context.Background(), watchlists.ItemsRequest{ListID: "wl-1"})
	if err != nil {
		t.Fatal(err)
	}
	if len(got.Items) != 1 || got.Items[0].Symbol != "MU" || got.Items[0].ObjectType != "equity" {
		t.Fatalf("%+v", got)
	}
	if s.LastName() != "get_watchlist_items" {
		t.Fatalf("%s", s.LastName())
	}
	want := map[string]any{"list_id": "wl-1"}
	if diff := cmp.Diff(want, s.LastArgs()); diff != "" {
		t.Fatal(diff)
	}
}

func TestLists_titleFallsBackToName(t *testing.T) {
	t.Parallel()
	c := watchlists.New(client.Func(func(ctx context.Context, name string, args map[string]any) (json.RawMessage, error) {
		return json.RawMessage(`{"watchlists":[{"id":"wl-1","title":"tech"},{"id":"wl-2","name":"Other"}]}`), nil
	}))
	got, err := c.Lists(context.Background(), watchlists.ListsRequest{})
	if err != nil {
		t.Fatal(err)
	}
	want := watchlists.ListsResult{
		Watchlists: []watchlists.Watchlist{
			{ID: "wl-1", Title: "tech"},
			{ID: "wl-2", Title: "Other"},
		},
	}
	if diff := cmp.Diff(want, got); diff != "" {
		t.Fatal(diff)
	}
}

func TestItems_keepsObjectTypesAndNestedSymbol(t *testing.T) {
	t.Parallel()
	c := watchlists.New(client.Func(func(ctx context.Context, name string, args map[string]any) (json.RawMessage, error) {
		return json.RawMessage(`{"items":[
			{"symbol":"MU","object_type":"equity"},
			{"symbol":"BTC","object_type":"crypto"},
			{"instrument":{"symbol":"SPY","type":"etf"}}
		]}`), nil
	}))
	got, err := c.Items(context.Background(), watchlists.ItemsRequest{ListID: "wl-1"})
	if err != nil {
		t.Fatal(err)
	}
	want := watchlists.ItemsResult{
		Items: []watchlists.Item{
			{Symbol: "MU", ObjectType: "equity"},
			{Symbol: "BTC", ObjectType: "crypto"},
			{
				Symbol:     "SPY",
				ObjectType: "etf",
				Instrument: &watchlists.Instrument{Symbol: "SPY", Type: "etf"},
			},
		},
	}
	if diff := cmp.Diff(want, got); diff != "" {
		t.Fatal(diff)
	}
}