robinhood-agentic-mcp/equity/write_test.go

package equity_test

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

	decimal "github.com/alpacahq/alpacadecimal"
	"github.com/google/go-cmp/cmp"
	"s1d3sw1ped/robinhood-agentic-mcp/client"
	"s1d3sw1ped/robinhood-agentic-mcp/equity"
	"s1d3sw1ped/robinhood-agentic-mcp/internal/rhntest"
)

func TestPlaceOrder_decimalStrings(t *testing.T) {
	t.Parallel()
	var got map[string]any
	c := equity.New(client.Func(func(ctx context.Context, name string, args map[string]any) (json.RawMessage, error) {
		if name != "place_equity_order" {
			t.Fatalf("%s", name)
		}
		got = args
		return json.RawMessage(`{"id":"o1"}`), nil
	}))
	qty := decimal.NewFromInt(3)
	px := decimal.RequireFromString("99.6")
	_, err := c.PlaceOrder(context.Background(), equity.PlaceOrderRequest{
		AccountNumber: "acct", Symbol: "MU", Side: client.Buy, Type: client.Limit,
		Qty: &qty, LimitPrice: &px, TimeInForce: client.GFD, RefID: "buy:2026-08-18:MU",
	})
	if err != nil {
		t.Fatal(err)
	}
	if got["type"] != "limit" || got["time_in_force"] != "gfd" || got["quantity"] != "3" {
		t.Fatalf("%+v", got)
	}
	if got["limit_price"] != "99.6" && got["limit_price"] != "99.60" {
		t.Fatalf("limit %v", got["limit_price"])
	}
	if got["ref_id"] != "buy:2026-08-18:MU" {
		t.Fatalf("ref %v", got["ref_id"])
	}
	if _, ok := got["idempotency_key"]; ok {
		t.Fatalf("idempotency_key %v", got["idempotency_key"])
	}
}

func TestPlaceOrder_ordersEnvelope(t *testing.T) {
	t.Parallel()
	c := equity.New(client.Func(func(ctx context.Context, name string, args map[string]any) (json.RawMessage, error) {
		return json.RawMessage(`{"orders":[{"id":"o2","symbol":"MU","side":"buy","quantity":"3","state":"confirmed"}]}`), nil
	}))
	got, err := c.PlaceOrder(context.Background(), equity.PlaceOrderRequest{Symbol: "MU"})
	if err != nil {
		t.Fatal(err)
	}
	if got.ID != "o2" || got.Symbol != "MU" || got.Side != "buy" || got.State != "confirmed" {
		t.Fatalf("%+v", got)
	}
	if !got.Qty.Equal(decimal.NewFromInt(3)) {
		t.Fatalf("qty %s", got.Qty)
	}
}

func TestEquity_writeArgs(t *testing.T) {
	t.Parallel()
	qty := decimal.NewFromInt(3)
	px := decimal.RequireFromString("99.6")
	dollars := decimal.RequireFromString("100")
	lotQty := decimal.RequireFromString("1.5")
	tests := []struct {
		name     string
		call     func(*equity.Client) error
		wantName string
		wantArgs map[string]any
	}{
		{
			name: "ReviewOrder",
			call: func(c *equity.Client) error {
				_, err := c.ReviewOrder(context.Background(), equity.PlaceOrderRequest{
					AccountNumber: "acct",
					Symbol:        "MU",
					Side:          client.Buy,
					Type:          client.Limit,
					Qty:           &qty,
					LimitPrice:    &px,
					TimeInForce:   client.GFD,
					MarketHours:   client.ExtendedHours,
					RefID:         "buy:2026-08-18:MU",
				})
				return err
			},
			wantName: "review_equity_order",
			wantArgs: map[string]any{
				"account_number": "acct",
				"symbol":         "MU",
				"side":           "buy",
				"type":           "limit",
				"quantity":       "3",
				"limit_price":    "99.6",
				"time_in_force":  "gfd",
				"market_hours":   "extended_hours",
			},
		},
		{
			name: "PlaceOrderStop",
			call: func(c *equity.Client) error {
				_, err := c.PlaceOrder(context.Background(), equity.PlaceOrderRequest{
					AccountNumber: "acct",
					Symbol:        "MU",
					Side:          client.Sell,
					Type:          client.Stop,
					Qty:           &qty,
					StopPrice:     &px,
					RefID:         "sell:2026-08-18:MU",
				})
				return err
			},
			wantName: "place_equity_order",
			wantArgs: map[string]any{
				"account_number": "acct",
				"symbol":         "MU",
				"side":           "sell",
				"type":           "stop_market",
				"quantity":       "3",
				"stop_price":     "99.6",
				"ref_id":         "sell:2026-08-18:MU",
			},
		},
		{
			name: "PlaceOrderDollarAndLots",
			call: func(c *equity.Client) error {
				_, err := c.PlaceOrder(context.Background(), equity.PlaceOrderRequest{
					AccountNumber: "acct",
					Symbol:        "MU",
					Side:          client.Sell,
					Type:          client.Market,
					DollarAmount:  &dollars,
					TaxLots: []equity.TaxLot{{
						OpenLotID: "lot-1",
						Quantity:  lotQty,
					}},
					RefID: "sell:lots",
				})
				return err
			},
			wantName: "place_equity_order",
			wantArgs: map[string]any{
				"account_number": "acct",
				"symbol":         "MU",
				"side":           "sell",
				"type":           "market",
				"dollar_amount":  "100",
				"tax_lots": []map[string]any{
					{"open_lot_id": "lot-1", "quantity": "1.5"},
				},
				"ref_id": "sell:lots",
			},
		},
		{
			name: "CancelOrder",
			call: func(c *equity.Client) error {
				return c.CancelOrder(context.Background(), equity.CancelOrderRequest{
					AccountNumber: "acct",
					OrderID:       "o1",
				})
			},
			wantName: "cancel_equity_order",
			wantArgs: map[string]any{"account_number": "acct", "order_id": "o1"},
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			var gotName string
			var gotArgs map[string]any
			c := equity.New(client.Func(func(ctx context.Context, name string, args map[string]any) (json.RawMessage, error) {
				gotName, gotArgs = name, args
				return json.RawMessage(`{"id":"o1"}`), 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)
			}
			if tc.name == "ReviewOrder" {
				if _, ok := gotArgs["ref_id"]; ok {
					t.Fatalf("ref_id on review: %+v", gotArgs)
				}
				if _, ok := gotArgs["idempotency_key"]; ok {
					t.Fatalf("idempotency_key on review: %+v", gotArgs)
				}
			}
			if tc.name == "PlaceOrderStop" {
				if _, ok := gotArgs["time_in_force"]; ok {
					t.Fatalf("time_in_force injected: %+v", gotArgs)
				}
			}
			if tc.wantName == "place_equity_order" {
				if _, ok := gotArgs["idempotency_key"]; ok {
					t.Fatalf("idempotency_key on place: %+v", gotArgs)
				}
			}
		})
	}
}

func TestPlaceOrder_rhntest(t *testing.T) {
	t.Parallel()
	s := rhntest.New(t)
	s.Set("place_equity_order", json.RawMessage(`{"id":"o1"}`))
	c := equity.New(&client.Client{URL: s.URL})
	qty := decimal.NewFromInt(3)
	px := decimal.RequireFromString("99.6")
	got, err := c.PlaceOrder(context.Background(), equity.PlaceOrderRequest{
		AccountNumber: "acct", Symbol: "MU", Side: client.Buy, Type: client.Limit,
		Qty: &qty, LimitPrice: &px, TimeInForce: client.GFD, RefID: "buy:2026-08-18:MU",
	})
	if err != nil {
		t.Fatal(err)
	}
	if got.ID != "o1" {
		t.Fatalf("%+v", got)
	}
	if s.LastName() != "place_equity_order" {
		t.Fatalf("%s", s.LastName())
	}
}

func TestReviewOrder_rhntest(t *testing.T) {
	t.Parallel()
	s := rhntest.New(t)
	s.Set("review_equity_order", json.RawMessage(`{"data":{"errors":["insufficient buying power"],"warnings":["PDT"]}}`))
	c := equity.New(&client.Client{URL: s.URL})
	qty := decimal.NewFromInt(3)
	got, err := c.ReviewOrder(context.Background(), equity.PlaceOrderRequest{
		AccountNumber: "acct", Symbol: "MU", Side: client.Buy, Type: client.Market, Qty: &qty,
	})
	if err != nil {
		t.Fatal(err)
	}
	if diff := cmp.Diff([]string{"insufficient buying power"}, got.Errors); diff != "" {
		t.Fatal(diff)
	}
	if diff := cmp.Diff([]string{"PDT"}, got.Warnings); diff != "" {
		t.Fatal(diff)
	}
	if s.LastName() != "review_equity_order" {
		t.Fatalf("%s", s.LastName())
	}
}

func TestCancelOrder_rhntest(t *testing.T) {
	t.Parallel()
	s := rhntest.New(t)
	s.Set("cancel_equity_order", json.RawMessage(`{}`))
	c := equity.New(&client.Client{URL: s.URL})
	if err := c.CancelOrder(context.Background(), equity.CancelOrderRequest{
		AccountNumber: "acct",
		OrderID:       "o1",
	}); err != nil {
		t.Fatal(err)
	}
	if s.LastName() != "cancel_equity_order" {
		t.Fatalf("%s", s.LastName())
	}
	want := map[string]any{"account_number": "acct", "order_id": "o1"}
	if diff := cmp.Diff(want, s.LastArgs()); diff != "" {
		t.Fatal(diff)
	}
}