Show correct message on NodeStatus failure

.gitignore

@@ -1,6 +1,8 @@
 /parts/
 /stage/
 /prime/
+.vscode/
+env/
 
 
 **.snap
@@ -20,4 +22,4 @@ k8s_*.txt
 /docs/tools/.sphinx/warnings.txt
 /docs/tools/.sphinx/.wordlist.dic
 /docs/tools/.sphinx/.doctrees/
-/docs/tools/.sphinx/node_modules
+/docs/tools/.sphinx/node_modules

src/k8s/cmd/k8s/k8s_config.go

@@ -32,11 +32,7 @@ func newKubeConfigCmd(env cmdutil.ExecutionEnvironment) *cobra.Command {
 				return
 			}
 
-			if _, err := client.NodeStatus(cmd.Context()); err != nil {
-				cmd.PrintErrln("Error: The node is not part of a Kubernetes cluster. You can bootstrap a new cluster with:\n\n  sudo k8s bootstrap")
-				env.Exit(1)
-				return
-			}
+			_ = GetNodeStatus(client, cmd, env)
 
 			ctx, cancel := context.WithTimeout(cmd.Context(), opts.timeout)
 			cobra.OnFinalize(cancel)

src/k8s/cmd/k8s/k8s_helm.go

@@ -24,11 +24,7 @@ func newHelmCmd(env cmdutil.ExecutionEnvironment) *cobra.Command {
 				return
 			}
 
-			if status, err := client.NodeStatus(cmd.Context()); err != nil {
-				cmd.PrintErrln("Error: The node is not part of a Kubernetes cluster. You can bootstrap a new cluster with:\n\n  sudo k8s bootstrap")
-				env.Exit(1)
-				return
-			} else if status.ClusterRole == apiv1.ClusterRoleWorker {
+			if status := GetNodeStatus(client, cmd, env); status.ClusterRole == apiv1.ClusterRoleWorker {
 				cmd.PrintErrln("Error: k8s helm commands are not allowed on worker nodes.")
 				env.Exit(1)
 				return

src/k8s/cmd/k8s/k8s_kubectl.go

@@ -24,11 +24,7 @@ func newKubectlCmd(env cmdutil.ExecutionEnvironment) *cobra.Command {
 				return
 			}
 
-			if status, err := client.NodeStatus(cmd.Context()); err != nil {
-				cmd.PrintErrln("Error: The node is not part of a Kubernetes cluster. You can bootstrap a new cluster with:\n\n  sudo k8s bootstrap")
-				env.Exit(1)
-				return
-			} else if status.ClusterRole == apiv1.ClusterRoleWorker {
+			if status := GetNodeStatus(client, cmd, env); status.ClusterRole == apiv1.ClusterRoleWorker {
 				cmd.PrintErrln("Error: k8s kubectl commands are not allowed on worker nodes.")
 				env.Exit(1)
 				return

src/k8s/cmd/k8s/k8s_local_node_status.go

@@ -22,14 +22,7 @@ func newLocalNodeStatusCommand(env cmdutil.ExecutionEnvironment) *cobra.Command
 				return
 			}
 
-			status, err := client.NodeStatus(cmd.Context())
-			if err != nil {
-				cmd.PrintErrf("Error: Failed to get the status of the local node.\n\nThe error was: %v\n", err)
-				env.Exit(1)
-				return
-			}
-
-			outputFormatter.Print(status)
+			outputFormatter.Print(GetNodeStatus(client, cmd, env))
 		},
 	}
 	cmd.Flags().StringVar(&opts.outputFormat, "output-format", "plain", "set the output format to one of plain, json or yaml")

src/k8s/cmd/k8s/k8s_status.go

@@ -35,11 +35,7 @@ func newStatusCmd(env cmdutil.ExecutionEnvironment) *cobra.Command {
 			ctx, cancel := context.WithTimeout(cmd.Context(), opts.timeout)
 			cobra.OnFinalize(cancel)
 
-			if _, err := client.NodeStatus(cmd.Context()); err != nil {
-				cmd.PrintErrln("Error: The node is not part of a Kubernetes cluster. You can bootstrap a new cluster with:\n\n  sudo k8s bootstrap")
-				env.Exit(1)
-				return
-			}
+			_ = GetNodeStatus(client, cmd, env)
 
 			status, err := client.ClusterStatus(ctx, opts.waitReady)
 			if err != nil {

src/k8s/cmd/k8s/node_status.go

@@ -0,0 +1,40 @@
+package k8s
+
+import (
+	"errors"
+	"net/http"
+
+	"github.com/spf13/cobra"
+
+	apiv1 "github.com/canonical/k8s/api/v1"
+	cmdutil "github.com/canonical/k8s/cmd/util"
+	"github.com/canonical/k8s/pkg/client/k8sd"
+
+	"github.com/canonical/lxd/shared/api"
+)
+
+// GetNodeStatus retrieves the NodeStatus from k8sd client.
+// If the daemon is not initialized, it exits with an error
+// describing that the cluster should be bootstrapped.
+// In case of any other errors it exits and shows the error.
+func GetNodeStatus(client k8sd.Client, cmd *cobra.Command, env cmdutil.ExecutionEnvironment) apiv1.NodeStatus {
+	status, err := client.NodeStatus(cmd.Context())
+	if err == nil {
+		return status
+	}
+
+	if errors.As(err, &api.StatusError{}) {
+		// the returned `ok` can be ignored since we're using errors.As()
+		// on the same type immediately before it
+		statusErr, _ := err.(api.StatusError)
+		if statusErr.Status() == http.StatusServiceUnavailable {
+			cmd.PrintErrln("Error: The node is not part of a Kubernetes cluster. You can bootstrap a new cluster with:\n\n  sudo k8s bootstrap")
+			env.Exit(1)
+			return status
+		}
+	}
+
+	cmd.PrintErrf("Error: Failed to retrieve the node status.\n\nThe error was: %v\n", err)
+	env.Exit(1)
+	return status
+}

src/k8s/cmd/k8s/node_status_test.go

@@ -0,0 +1,101 @@
+package k8s_test
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"net/http"
+	"testing"
+
+	apiv1 "github.com/canonical/k8s/api/v1"
+	"github.com/canonical/k8s/cmd/k8s"
+	cmdutil "github.com/canonical/k8s/cmd/util"
+	k8sdmock "github.com/canonical/k8s/pkg/client/k8sd/mock"
+	snapmock "github.com/canonical/k8s/pkg/snap/mock"
+	"github.com/canonical/lxd/shared/api"
+
+	. "github.com/onsi/gomega"
+)
+
+func TestGetNodeStatusUtil(t *testing.T) {
+	type testcase struct {
+		expectedCode     int
+		name             string
+		nodeStatusErr    error
+		expectedInStdErr []string
+		nodeStatusResult apiv1.NodeStatus
+	}
+
+	tests := []testcase{
+		{
+			name: "NoError",
+			nodeStatusResult: apiv1.NodeStatus{
+				Name: "name", Address: "addr",
+				ClusterRole: apiv1.ClusterRoleControlPlane, DatastoreRole: apiv1.DatastoreRoleVoter,
+			},
+		},
+		{
+			name:             "DaemonNotInitialized",
+			nodeStatusErr:    api.StatusErrorf(http.StatusServiceUnavailable, "Daemon not yet initialized"),
+			expectedCode:     1,
+			expectedInStdErr: []string{"The node is not part of a Kubernetes cluster. You can bootstrap a new cluster"},
+		},
+		{
+			name:             "RandomError",
+			nodeStatusErr:    errors.New("something went bad"),
+			expectedCode:     1,
+			expectedInStdErr: []string{"something went bad", "Failed to retrieve the node status"},
+		},
+		{
+			name:             "ContextCanceled",
+			nodeStatusErr:    context.Canceled,
+			expectedCode:     1,
+			expectedInStdErr: []string{context.Canceled.Error(), "Failed to retrieve the node status"},
+		},
+		{
+			name:             "ContextDeadlineExceeded",
+			nodeStatusErr:    context.DeadlineExceeded,
+			expectedCode:     1,
+			expectedInStdErr: []string{context.DeadlineExceeded.Error(), "Failed to retrieve the node status"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			g := NewWithT(t)
+
+			var (
+				returnCode int
+				stdout     = &bytes.Buffer{}
+				stderr     = &bytes.Buffer{}
+				mockClient = &k8sdmock.Mock{
+					NodeStatusErr:    tt.nodeStatusErr,
+					NodeStatusResult: tt.nodeStatusResult,
+				}
+				env = cmdutil.ExecutionEnvironment{
+					Stdout: stdout,
+					Stderr: stderr,
+					Getuid: func() int { return 0 },
+					Snap: &snapmock.Snap{
+						Mock: snapmock.Mock{
+							K8sdClient: mockClient,
+						},
+					},
+					Exit: func(rc int) { returnCode = rc },
+				}
+				cmd = k8s.NewRootCmd(env)
+			)
+
+			status := k8s.GetNodeStatus(mockClient, cmd, env)
+
+			g.Expect(returnCode).To(Equal(tt.expectedCode))
+			for _, exp := range tt.expectedInStdErr {
+				g.Expect(stderr.String()).To(ContainSubstring(exp))
+			}
+
+			if tt.expectedCode == 0 {
+				g.Expect(status).To(Equal(tt.nodeStatusResult))
+			}
+		})
+	}
+}

src/k8s/pkg/k8sd/api/node.go

@@ -14,7 +14,7 @@ func (e *Endpoints) getNodeStatus(s *state.State, r *http.Request) response.Resp
 
 	status, err := impl.GetLocalNodeStatus(r.Context(), s, snap)
 	if err != nil {
-		response.InternalError(err)
+		return response.InternalError(err)
 	}
 
 	result := apiv1.GetNodeStatusResponse{