clean up functions

Author: patelspratik

pkg/cmd/grantssh/grantssh.go

@@ -139,17 +139,7 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt
 	// Capture the target node
 	var node *nodev1.ExternalNode
 	if opts.interactive {
-		resp, listErr := client.ListNodes(ctx, connect.NewRequest(&nodev1.ListNodesRequest{
-			OrganizationId: org.ID,
-		}))
-		if listErr != nil {
-			return breverrors.WrapAndTrace(listErr)
-		}
-		nodes := resp.Msg.GetItems()
-		if len(nodes) == 0 {
-			return fmt.Errorf("no nodes found in organization")
-		}
-		node, err = register.SelectNodeFromList(ctx, t, deps.prompter, deps.registrationStore, nodes)
+		node, err = register.SelectNodeFromList(ctx, org.ID, t, deps.prompter, deps.registrationStore, client)
 	} else {
 		node, err = helpers.ResolveNodeByName(ctx, client, org.ID, opts.nodeName)
 	}

pkg/cmd/register/sshkeys.go

@@ -12,8 +12,10 @@ import (
 	"strings"
 	"time"
 
+	nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect"
 	nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
 	"connectrpc.com/connect"
+	breverrors "github.com/brevdev/brev-cli/pkg/errors"
 
 	"github.com/brevdev/brev-cli/pkg/config"
 	"github.com/brevdev/brev-cli/pkg/entity"
@@ -25,7 +27,14 @@ import (
 // SelectNodeFromList prompts the user to select a node from the list, then fetches and returns the full node.
 // When this machine is registered, that node is marked " — this node" in the list.
 // Call this after ListNodes when in interactive mode (same pattern as SelectOrganizationInteractive).
-func SelectNodeFromList(ctx context.Context, t *terminal.Terminal, prompter terminal.Selector, registrationStore RegistrationStore, nodes []*nodev1.ExternalNode) (*nodev1.ExternalNode, error) {
+func SelectNodeFromList(ctx context.Context, orgID string, t *terminal.Terminal, prompter terminal.Selector, registrationStore RegistrationStore, client nodev1connect.ExternalNodeServiceClient) (*nodev1.ExternalNode, error) {
+	resp, listErr := client.ListNodes(ctx, connect.NewRequest(&nodev1.ListNodesRequest{
+		OrganizationId: orgID,
+	}))
+	if listErr != nil {
+		return nil, breverrors.WrapAndTrace(listErr)
+	}
+	nodes := resp.Msg.GetItems()
 	if len(nodes) == 0 {
 		return nil, fmt.Errorf("no nodes found in organization")
 	}

pkg/cmd/register/sshkeys_test.go

@@ -8,7 +8,9 @@ import (
 	"strings"
 	"testing"
 
+	nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect"
 	nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
+	"connectrpc.com/connect"
 
 	"github.com/brevdev/brev-cli/pkg/terminal"
 )
@@ -249,6 +251,17 @@ func (s *captureSelector) Select(_ string, items []string) string {
 	return items[s.returnIdx]
 }
 
+// fakeNodeClient implements nodev1connect.ExternalNodeServiceClient, returning
+// canned nodes from ListNodes
+type fakeNodeClient struct {
+	nodev1connect.ExternalNodeServiceClient
+	nodes []*nodev1.ExternalNode
+}
+
+func (f *fakeNodeClient) ListNodes(_ context.Context, _ *connect.Request[nodev1.ListNodesRequest]) (*connect.Response[nodev1.ListNodesResponse], error) {
+	return connect.NewResponse(&nodev1.ListNodesResponse{Items: f.nodes}), nil
+}
+
 func makeNodes(ids ...string) []*nodev1.ExternalNode {
 	nodes := make([]*nodev1.ExternalNode, len(ids))
 	for i, id := range ids {
@@ -265,9 +278,9 @@ func TestSelectNodeFromList_ThisNodeMovedToFront(t *testing.T) {
 	regStore := &mockRegistrationStore{
 		reg: &DeviceRegistration{ExternalNodeID: "c"},
 	}
-	nodes := makeNodes("a", "b", "c", "d")
+	client := &fakeNodeClient{nodes: makeNodes("a", "b", "c", "d")}
 
-	got, err := SelectNodeFromList(context.Background(), terminal.New(), sel, regStore, nodes)
+	got, err := SelectNodeFromList(context.Background(), "org1", terminal.New(), sel, regStore, client)
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -293,9 +306,9 @@ func TestSelectNodeFromList_ThisNodeAlreadyFirst(t *testing.T) {
 	regStore := &mockRegistrationStore{
 		reg: &DeviceRegistration{ExternalNodeID: "a"},
 	}
-	nodes := makeNodes("a", "b", "c")
+	client := &fakeNodeClient{nodes: makeNodes("a", "b", "c")}
 
-	got, err := SelectNodeFromList(context.Background(), terminal.New(), sel, regStore, nodes)
+	got, err := SelectNodeFromList(context.Background(), "org1", terminal.New(), sel, regStore, client)
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -311,10 +324,9 @@ func TestSelectNodeFromList_ThisNodeAlreadyFirst(t *testing.T) {
 func TestSelectNodeFromList_NoRegistration(t *testing.T) {
 	sel := &captureSelector{returnIdx: 1}
 	regStore := &mockRegistrationStore{} // no registration
+	client := &fakeNodeClient{nodes: makeNodes("a", "b", "c")}
 
-	nodes := makeNodes("a", "b", "c")
-
-	got, err := SelectNodeFromList(context.Background(), terminal.New(), sel, regStore, nodes)
+	got, err := SelectNodeFromList(context.Background(), "org1", terminal.New(), sel, regStore, client)
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}

pkg/cmd/revokessh/revokessh.go

@@ -128,17 +128,7 @@ func runRevokeSSH(ctx context.Context, t *terminal.Terminal, s RevokeSSHStore, o
 	// Capture the target node
 	var selectedNode *nodev1.ExternalNode
 	if opts.interactive {
-		resp, listErr := client.ListNodes(ctx, connect.NewRequest(&nodev1.ListNodesRequest{
-			OrganizationId: selectedOrg.ID,
-		}))
-		if listErr != nil {
-			return breverrors.WrapAndTrace(listErr)
-		}
-		nodes := resp.Msg.GetItems()
-		if len(nodes) == 0 {
-			return fmt.Errorf("no nodes found in organization")
-		}
-		node, err := register.SelectNodeFromList(ctx, t, deps.prompter, deps.registrationStore, nodes)
+		node, err := register.SelectNodeFromList(ctx, selectedOrg.ID, t, deps.prompter, deps.registrationStore, client)
 		if err != nil {
 			return err //nolint:wrapcheck // do not present stack trace for this error
 		}