Use SDK HTTP client for network checks, return error on check failure

Author: simonfaltum

cmd/doctor/checks.go

@@ -1,6 +1,8 @@
 package doctor
 
 import (
+	"context"
+	"crypto/tls"
 	"errors"
 	"fmt"
 	"io"
@@ -11,6 +13,7 @@ import (
 	"github.com/databricks/cli/internal/build"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/spf13/cobra"
@@ -230,14 +233,16 @@ func checkNetwork(cmd *cobra.Command, cfg *config.Config, resolveErr error, w *d
 	}
 
 	if w != nil {
-		return checkNetworkWithHost(cmd, w.Config.Host)
+		return checkNetworkWithHost(cmd, w.Config.Host, configuredNetworkHTTPClient(w.Config))
 	}
 
-	return checkNetworkWithHost(cmd, cfg.Host)
+	log.Warnf(cmd.Context(), "workspace client unavailable for network check, falling back to default HTTP client")
+	return checkNetworkWithHost(cmd, cfg.Host, http.DefaultClient)
 }
 
-func checkNetworkWithHost(cmd *cobra.Command, host string) CheckResult {
-	ctx := cmd.Context()
+func checkNetworkWithHost(cmd *cobra.Command, host string, client *http.Client) CheckResult {
+	ctx, cancel := context.WithTimeout(cmd.Context(), networkTimeout)
+	defer cancel()
 
 	if host == "" {
 		return CheckResult{
@@ -257,7 +262,6 @@ func checkNetworkWithHost(cmd *cobra.Command, host string) CheckResult {
 		}
 	}
 
-	client := &http.Client{Timeout: networkTimeout}
 	resp, err := client.Do(req)
 	if err != nil {
 		return CheckResult{
@@ -276,3 +280,33 @@ func checkNetworkWithHost(cmd *cobra.Command, host string) CheckResult {
 		Message: host + " is reachable",
 	}
 }
+
+func configuredNetworkHTTPClient(cfg *config.Config) *http.Client {
+	return &http.Client{
+		Transport: configuredNetworkHTTPTransport(cfg),
+	}
+}
+
+func configuredNetworkHTTPTransport(cfg *config.Config) http.RoundTripper {
+	if cfg.HTTPTransport != nil {
+		return cfg.HTTPTransport
+	}
+
+	if !cfg.InsecureSkipVerify {
+		return http.DefaultTransport
+	}
+
+	transport, ok := http.DefaultTransport.(*http.Transport)
+	if !ok {
+		return http.DefaultTransport
+	}
+
+	clone := transport.Clone()
+	if clone.TLSClientConfig != nil {
+		clone.TLSClientConfig = clone.TLSClientConfig.Clone()
+	} else {
+		clone.TLSClientConfig = &tls.Config{}
+	}
+	clone.TLSClientConfig.InsecureSkipVerify = true
+	return clone
+}

cmd/doctor/doctor.go

@@ -2,6 +2,7 @@ package doctor
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 
@@ -22,10 +23,12 @@ type CheckResult struct {
 // New returns the doctor command.
 func New() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "doctor",
-		Args:    root.NoArgs,
-		Short:   "Validate your Databricks CLI setup",
-		GroupID: "development",
+		Use:           "doctor",
+		Args:          root.NoArgs,
+		Short:         "Validate your Databricks CLI setup",
+		GroupID:       "development",
+		SilenceUsage:  true,
+		SilenceErrors: true,
 	}
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
@@ -39,13 +42,19 @@ func New() *cobra.Command {
 			}
 			buf = append(buf, '\n')
 			_, err = cmd.OutOrStdout().Write(buf)
-			return err
+			if err != nil {
+				return err
+			}
 		case flags.OutputText:
 			renderResults(cmd.OutOrStdout(), results)
-			return nil
 		default:
 			return fmt.Errorf("unknown output type %s", root.OutputType(cmd))
 		}
+
+		if hasFailedChecks(results) {
+			return errors.New("one or more checks failed")
+		}
+		return nil
 	}
 
 	return cmd
@@ -77,3 +86,12 @@ func renderResults(w io.Writer, results []CheckResult) {
 		fmt.Fprintln(w, msg)
 	}
 }
+
+func hasFailedChecks(results []CheckResult) bool {
+	for _, result := range results {
+		if result.Status == statusFail {
+			return true
+		}
+	}
+	return false
+}

cmd/doctor/doctor_test.go

@@ -60,6 +60,12 @@ func (m *noConfigProfiler) GetPath(_ context.Context) (string, error) {
 	return m.path, nil
 }
 
+type roundTripFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
+	return f(r)
+}
+
 func newTestCmd(ctx context.Context) *cobra.Command {
 	cmd := &cobra.Command{}
 	cmd.SetContext(ctx)
@@ -285,7 +291,7 @@ func TestCheckNetworkReachable(t *testing.T) {
 	ctx := cmdio.MockDiscard(t.Context())
 	cmd := newTestCmd(ctx)
 
-	result := checkNetworkWithHost(cmd, srv.URL)
+	result := checkNetworkWithHost(cmd, srv.URL, http.DefaultClient)
 	assert.Equal(t, "Network", result.Name)
 	assert.Equal(t, statusPass, result.Status)
 	assert.Contains(t, result.Message, "reachable")
@@ -295,21 +301,26 @@ func TestCheckNetworkNoHost(t *testing.T) {
 	ctx := cmdio.MockDiscard(t.Context())
 	cmd := newTestCmd(ctx)
 
-	result := checkNetworkWithHost(cmd, "")
+	result := checkNetworkWithHost(cmd, "", http.DefaultClient)
 	assert.Equal(t, "Network", result.Name)
 	assert.Equal(t, statusFail, result.Status)
 	assert.Contains(t, result.Message, "No host configured")
 }
 
-func TestCheckNetworkWithClient(t *testing.T) {
-	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.WriteHeader(http.StatusOK)
-	}))
-	defer srv.Close()
-
+func TestCheckNetworkUsesWorkspaceClientTransport(t *testing.T) {
 	w, err := databricks.NewWorkspaceClient((*databricks.Config)(&config.Config{
-		Host:  srv.URL,
+		Host:  "https://example.com",
 		Token: "test-token",
+		HTTPTransport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+			assert.Equal(t, http.MethodHead, r.Method)
+			assert.Equal(t, "https://example.com", r.URL.String())
+			return &http.Response{
+				StatusCode: http.StatusOK,
+				Body:       http.NoBody,
+				Header:     make(http.Header),
+				Request:    r,
+			}, nil
+		}),
 	}))
 	require.NoError(t, err)
 
@@ -390,6 +401,38 @@ func TestRenderResultsJSONOmitsEmptyDetail(t *testing.T) {
 	assert.NotContains(t, string(buf), "detail")
 }
 
+func TestHasFailedChecks(t *testing.T) {
+	tests := []struct {
+		name    string
+		results []CheckResult
+		want    bool
+	}{
+		{
+			name: "no failures",
+			results: []CheckResult{
+				{Name: "Test", Status: statusPass},
+				{Name: "Info", Status: statusInfo},
+				{Name: "Warn", Status: statusWarn},
+			},
+			want: false,
+		},
+		{
+			name: "has failure",
+			results: []CheckResult{
+				{Name: "Test", Status: statusPass},
+				{Name: "Broken", Status: statusFail},
+			},
+			want: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, hasFailedChecks(tt.results))
+		})
+	}
+}
+
 func TestNewCommandJSON(t *testing.T) {
 	clearConfigEnv(t)
 
@@ -412,7 +455,7 @@ func TestNewCommandJSON(t *testing.T) {
 	cmd.SetArgs([]string{"--output", "json"})
 
 	err := cmd.Execute()
-	require.NoError(t, err)
+	require.ErrorContains(t, err, "one or more checks failed")
 
 	var results []CheckResult
 	err = json.Unmarshal(buf.Bytes(), &results)
@@ -445,7 +488,7 @@ func TestNewCommandJSONTrailingNewline(t *testing.T) {
 	cmd.SetArgs([]string{"--output", "json"})
 
 	err := cmd.Execute()
-	require.NoError(t, err)
+	require.ErrorContains(t, err, "one or more checks failed")
 	assert.Positive(t, buf.Len())
 	assert.Equal(t, byte('\n'), buf.Bytes()[buf.Len()-1])
 }