retry error

Author: drewmalin

pkg/cmd/register/register_test.go

@@ -499,3 +499,125 @@ Peers count: 0/0 Connected`
 		})
 	}
 }
+
+func TestIsSSHConnectionError(t *testing.T) {
+	t.Run("nil", func(t *testing.T) {
+		if IsSSHConnectionError(nil) {
+			t.Error("IsSSHConnectionError(nil) should be false")
+		}
+	})
+	t.Run("plain_error", func(t *testing.T) {
+		if IsSSHConnectionError(fmt.Errorf("some error")) {
+			t.Error("IsSSHConnectionError(plain error) should be false")
+		}
+	})
+	t.Run("connection_error_type", func(t *testing.T) {
+		err := &sshConnectionError{err: fmt.Errorf("transient")}
+		if !IsSSHConnectionError(err) {
+			t.Error("IsSSHConnectionError(sshConnectionError) should be true")
+		}
+	})
+	t.Run("wrapped_connection_error", func(t *testing.T) {
+		err := fmt.Errorf("wrapped: %w", &sshConnectionError{err: fmt.Errorf("transient")})
+		if !IsSSHConnectionError(err) {
+			t.Error("IsSSHConnectionError(wrapped sshConnectionError) should be true")
+		}
+	})
+}
+
+func Test_runRegister_GrantSSH_retries_on_connection_error_then_succeeds(t *testing.T) {
+	regStore := &mockRegistrationStore{}
+
+	store := &mockRegisterStore{
+		user:  &entity.User{ID: "user_1"},
+		org:   &entity.Organization{ID: "org_123", Name: "TestOrg"},
+		home:  "/home/testuser/.brev",
+		token: "tok",
+	}
+
+	var grantCalls int
+	svc := &fakeNodeService{
+		addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) {
+			return &nodev1.AddNodeResponse{
+				ExternalNode: &nodev1.ExternalNode{
+					ExternalNodeId: "unode_abc",
+					OrganizationId: "org_123",
+					Name:           req.GetName(),
+					DeviceId:       req.GetDeviceId(),
+					ConnectivityInfo: &nodev1.ConnectivityInfo{
+						RegistrationCommand: "netbird up --key abc",
+					},
+				},
+			}, nil
+		},
+		grantNodeSSHAccessFn: func(_ *nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) {
+			grantCalls++
+			if grantCalls < 2 {
+				return nil, connect.NewError(connect.CodeInternal, nil)
+			}
+			return &nodev1.GrantNodeSSHAccessResponse{}, nil
+		},
+	}
+
+	deps, server := testRegisterDeps(t, svc, regStore)
+	defer server.Close()
+
+	deps.prompter = mockConfirmer{confirm: true}
+
+	term := terminal.New()
+	err := runRegister(context.Background(), term, store, "My Spark", deps)
+	if err != nil {
+		t.Fatalf("runRegister failed: %v", err)
+	}
+
+	if grantCalls != 2 {
+		t.Errorf("expected GrantNodeSSHAccess to be called 2 times (retry once), got %d", grantCalls)
+	}
+}
+
+func Test_runRegister_GrantSSH_no_retry_on_permanent_error(t *testing.T) {
+	regStore := &mockRegistrationStore{}
+
+	store := &mockRegisterStore{
+		user:  &entity.User{ID: "user_1"},
+		org:   &entity.Organization{ID: "org_123", Name: "TestOrg"},
+		home:  "/home/testuser/.brev",
+		token: "tok",
+	}
+
+	var grantCalls int
+	svc := &fakeNodeService{
+		addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) {
+			return &nodev1.AddNodeResponse{
+				ExternalNode: &nodev1.ExternalNode{
+					ExternalNodeId: "unode_abc",
+					OrganizationId: "org_123",
+					Name:           req.GetName(),
+					DeviceId:       req.GetDeviceId(),
+					ConnectivityInfo: &nodev1.ConnectivityInfo{
+						RegistrationCommand: "netbird up --key abc",
+					},
+				},
+			}, nil
+		},
+		grantNodeSSHAccessFn: func(_ *nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) {
+			grantCalls++
+			return nil, connect.NewError(connect.CodePermissionDenied, nil)
+		},
+	}
+
+	deps, server := testRegisterDeps(t, svc, regStore)
+	defer server.Close()
+
+	deps.prompter = mockConfirmer{confirm: true}
+
+	term := terminal.New()
+	err := runRegister(context.Background(), term, store, "My Spark", deps)
+	if err != nil {
+		t.Fatalf("runRegister should not fail the overall flow when SSH grant fails: %v", err)
+	}
+
+	if grantCalls != 1 {
+		t.Errorf("expected GrantNodeSSHAccess to be called once (no retry on permanent error), got %d", grantCalls)
+	}
+}

pkg/cmd/register/rpcclient_test.go

@@ -143,9 +143,10 @@ func Test_toProtoNodeSpec_MinimalFields(t *testing.T) {
 // fakeNodeService implements the server side of ExternalNodeService for testing.
 type fakeNodeService struct {
 	nodev1connect.UnimplementedExternalNodeServiceHandler
-	addNodeFn    func(*nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error)
-	removeNodeFn func(*nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error)
-	getNodeFn    func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error)
+	addNodeFn            func(*nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error)
+	removeNodeFn         func(*nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error)
+	getNodeFn            func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error)
+	grantNodeSSHAccessFn func(*nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error)
 }
 
 func (f *fakeNodeService) AddNode(_ context.Context, req *connect.Request[nodev1.AddNodeRequest]) (*connect.Response[nodev1.AddNodeResponse], error) {
@@ -175,6 +176,17 @@ func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1
 	return connect.NewResponse(resp), nil
 }
 
+func (f *fakeNodeService) GrantNodeSSHAccess(_ context.Context, req *connect.Request[nodev1.GrantNodeSSHAccessRequest]) (*connect.Response[nodev1.GrantNodeSSHAccessResponse], error) {
+	if f.grantNodeSSHAccessFn == nil {
+		return nil, connect.NewError(connect.CodeUnimplemented, nil)
+	}
+	resp, err := f.grantNodeSSHAccessFn(req.Msg)
+	if err != nil {
+		return nil, err
+	}
+	return connect.NewResponse(resp), nil
+}
+
 func Test_NewNodeServiceClient_AddNode(t *testing.T) {
 	svc := &fakeNodeService{
 		addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) {