Add mocks and UTs for docker model status

Makefile

@@ -1,4 +1,4 @@
-.PHONY: all build clean link
+.PHONY: all build clean link mock unit-tests
 
 BINARY_NAME=model-cli
 
@@ -35,6 +35,17 @@ release:
 	GOOS=windows GOARCH=amd64 go build -ldflags="-s -w -X github.com/docker/model-cli/commands.Version=$(VERSION)" -o dist/windows-amd64/$(PLUGIN_NAME).exe .
 	@echo "Release build complete: $(PLUGIN_NAME) version '$(VERSION)'"
 
+mock:
+	@echo "Generating mocks..."
+	@mkdir -p mocks
+	@go generate ./...
+	@echo "Mocks generated!"
+
+unit-tests:
+	@echo "Running unit tests..."
+	@go test -v ./...
+	@echo "Unit tests completed!"
+
 clean:
 	@echo "Cleaning up..."
 	@rm -f $(BINARY_NAME)

commands/inspect.go

@@ -7,7 +7,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func newInspectCmd() *cobra.Command {
+func newInspectCmd(desktopClient *desktop.Client) *cobra.Command {
 	var openai bool
 	c := &cobra.Command{
 		Use:   "inspect MODEL",
@@ -24,11 +24,7 @@ func newInspectCmd() *cobra.Command {
 		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			model := args[0]
-			client, err := desktop.New()
-			if err != nil {
-				return fmt.Errorf("Failed to create Docker client: %v\n", err)
-			}
-			model, err = client.List(false, openai, model)
+			model, err := desktopClient.List(false, openai, model)
 			if err != nil {
 				err = handleClientError(err, "Failed to list models")
 				return handleNotRunningError(err)

commands/list.go

@@ -1,24 +1,18 @@
 package commands
 
 import (
-	"fmt"
-
 	"github.com/docker/model-cli/desktop"
 	"github.com/spf13/cobra"
 )
 
-func newListCmd() *cobra.Command {
+func newListCmd(desktopClient *desktop.Client) *cobra.Command {
 	var jsonFormat, openai bool
 	c := &cobra.Command{
 		Use:     "list",
 		Aliases: []string{"ls"},
 		Short:   "List the available models that can be run with the Docker Model Runner",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			client, err := desktop.New()
-			if err != nil {
-				return fmt.Errorf("Failed to create Docker client: %v\n", err)
-			}
-			models, err := client.List(jsonFormat, openai, "")
+			models, err := desktopClient.List(jsonFormat, openai, "")
 			if err != nil {
 				err = handleClientError(err, "Failed to list models")
 				return handleNotRunningError(err)

commands/pull.go

@@ -7,7 +7,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func newPullCmd() *cobra.Command {
+func newPullCmd(desktopClient *desktop.Client) *cobra.Command {
 	c := &cobra.Command{
 		Use:   "pull MODEL",
 		Short: "Download a model",
@@ -23,11 +23,7 @@ func newPullCmd() *cobra.Command {
 		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			model := args[0]
-			client, err := desktop.New()
-			if err != nil {
-				return fmt.Errorf("Failed to create Docker client: %v\n", err)
-			}
-			response, err := client.Pull(model)
+			response, err := desktopClient.Pull(model)
 			if err != nil {
 				err = handleClientError(err, "Failed to pull model")
 				return handleNotRunningError(err)

commands/rm.go

@@ -7,7 +7,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func newRemoveCmd() *cobra.Command {
+func newRemoveCmd(desktopClient *desktop.Client) *cobra.Command {
 	c := &cobra.Command{
 		Use:   "rm MODEL",
 		Short: "Remove a model downloaded from Docker Hub",
@@ -23,11 +23,7 @@ func newRemoveCmd() *cobra.Command {
 		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			model := args[0]
-			client, err := desktop.New()
-			if err != nil {
-				return fmt.Errorf("Failed to create Docker client: %v\n", err)
-			}
-			response, err := client.Remove(model)
+			response, err := desktopClient.Remove(model)
 			if err != nil {
 				err = handleClientError(err, "Failed to remove model")
 				return handleNotRunningError(err)

commands/root.go

@@ -1,20 +1,38 @@
 package commands
 
-import "github.com/spf13/cobra"
+import (
+	"fmt"
+	"os"
+
+	"github.com/docker/docker/client"
+	"github.com/docker/model-cli/desktop"
+	"github.com/docker/pinata/common/pkg/engine"
+	"github.com/docker/pinata/common/pkg/paths"
+	"github.com/spf13/cobra"
+)
 
 func NewRootCmd() *cobra.Command {
 	rootCmd := &cobra.Command{
 		Use:   "model",
 		Short: "Docker Model Runner",
 	}
+	dockerClient, err := client.NewClientWithOpts(
+		// TODO: Make sure it works while running in Windows containers mode.
+		client.WithHost(paths.HostServiceSockets().DockerHost(engine.Linux)),
+	)
+	if err != nil {
+		fmt.Println("Failed to create Docker client:", err)
+		os.Exit(1)
+	}
+	desktopClient := desktop.New(dockerClient.HTTPClient())
 	rootCmd.AddCommand(
 		newVersionCmd(),
-		newStatusCmd(),
-		newPullCmd(),
-		newListCmd(),
-		newRunCmd(),
-		newRemoveCmd(),
-		newInspectCmd(),
+		newStatusCmd(desktopClient),
+		newPullCmd(desktopClient),
+		newListCmd(desktopClient),
+		newRunCmd(desktopClient),
+		newRemoveCmd(desktopClient),
+		newInspectCmd(desktopClient),
 	)
 	return rootCmd
 }

commands/run.go

@@ -11,7 +11,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func newRunCmd() *cobra.Command {
+func newRunCmd(desktopClient *desktop.Client) *cobra.Command {
 	var debug bool
 
 	cmdArgs := "MODEL [PROMPT]"
@@ -32,25 +32,20 @@ func newRunCmd() *cobra.Command {
 				}
 			}
 
-			client, err := desktop.New()
-			if err != nil {
-				return fmt.Errorf("Failed to create Docker client: %v\n", err)
-			}
-
-			if _, err := client.List(false, false, model); err != nil {
+			if _, err := desktopClient.List(false, false, model); err != nil {
 				if !errors.Is(err, desktop.ErrNotFound) {
 					return handleNotRunningError(handleClientError(err, "Failed to list models"))
 				}
 				cmd.Println("Unable to find model '" + model + "' locally. Pulling from the server.")
-				response, err := client.Pull(model)
+				response, err := desktopClient.Pull(model)
 				if err != nil {
 					return handleNotRunningError(handleClientError(err, "Failed to pull model"))
 				}
 				cmd.Println(response)
 			}
 
 			if prompt != "" {
-				if err := client.Chat(model, prompt); err != nil {
+				if err := desktopClient.Chat(model, prompt); err != nil {
 					return handleClientError(err, "Failed to generate a response")
 				}
 				cmd.Println()
@@ -74,7 +69,7 @@ func newRunCmd() *cobra.Command {
 					continue
 				}
 
-				if err := client.Chat(model, userInput); err != nil {
+				if err := desktopClient.Chat(model, userInput); err != nil {
 					cmd.PrintErr(handleClientError(err, "Failed to generate a response"))
 					cmd.Print("> ")
 					continue

commands/status.go

@@ -1,37 +1,34 @@
 package commands
 
 import (
-	"fmt"
 	"os"
 
 	"github.com/docker/cli/cli-plugins/hooks"
 	"github.com/docker/model-cli/desktop"
 	"github.com/spf13/cobra"
 )
 
-func newStatusCmd() *cobra.Command {
+func newStatusCmd(desktopClient *desktop.Client) *cobra.Command {
 	c := &cobra.Command{
 		Use:   "status",
 		Short: "Check if the Docker Model Runner is running",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			client, err := desktop.New()
-			if err != nil {
-				return fmt.Errorf("Failed to create Docker client: %v\n", err)
-			}
-			status := client.Status()
+			status := desktopClient.Status()
 			if status.Error != nil {
-				return fmt.Errorf("Failed to get Docker Model Runner status: %v\n", err)
+				return handleClientError(status.Error, "Failed to get Docker Model Runner status")
 			}
 			if status.Running {
 				cmd.Println("Docker Model Runner is running")
 			} else {
 				cmd.Println("Docker Model Runner is not running")
-				hooks.PrintNextSteps(os.Stdout, []string{enableViaCLI, enableViaGUI})
-				os.Exit(1)
+				hooks.PrintNextSteps(cmd.OutOrStdout(), []string{enableViaCLI, enableViaGUI})
+				osExit(1)
 			}
 
 			return nil
 		},
 	}
 	return c
 }
+
+var osExit = os.Exit

commands/status_test.go

@@ -0,0 +1,113 @@
+package commands
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/docker/cli/cli-plugins/hooks"
+	"github.com/docker/model-cli/desktop"
+	mockdesktop "github.com/docker/model-cli/mocks"
+	"github.com/docker/pinata/common/pkg/inference"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+)
+
+func TestStatus(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockBody := io.NopCloser(strings.NewReader(""))
+
+	tests := []struct {
+		name           string
+		doResponse     *http.Response
+		doErr          error
+		expectExit     bool
+		expectedErr    error
+		expectedOutput string
+	}{
+		{
+			name:           "running",
+			doResponse:     &http.Response{StatusCode: http.StatusOK, Body: mockBody},
+			doErr:          nil,
+			expectExit:     false,
+			expectedErr:    nil,
+			expectedOutput: "Docker Model Runner is running\n",
+		},
+		{
+			name:        "not running",
+			doResponse:  &http.Response{StatusCode: http.StatusServiceUnavailable, Body: mockBody},
+			doErr:       nil,
+			expectExit:  true,
+			expectedErr: nil,
+			expectedOutput: func() string {
+				buf := new(bytes.Buffer)
+				fmt.Fprintln(buf, "Docker Model Runner is not running")
+				hooks.PrintNextSteps(buf, []string{enableViaCLI, enableViaGUI})
+				return buf.String()
+			}(),
+		},
+		{
+			name:       "request with error",
+			doResponse: &http.Response{StatusCode: http.StatusInternalServerError, Body: mockBody},
+			doErr:      nil,
+			expectExit: false,
+			expectedErr: handleClientError(
+				fmt.Errorf("unexpected status code: %d", http.StatusInternalServerError),
+				"Failed to get Docker Model Runner status",
+			),
+			expectedOutput: "",
+		},
+		{
+			name:       "failed request",
+			doResponse: nil,
+			doErr:      fmt.Errorf("failed to make request"),
+			expectExit: false,
+			expectedErr: handleClientError(
+				fmt.Errorf("error querying %s: %w", inference.ModelsPrefix, fmt.Errorf("failed to make request")),
+				"Failed to get Docker Model Runner status",
+			),
+			expectedOutput: "",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			client := mockdesktop.NewMockDockerHttpClient(ctrl)
+			req, err := http.NewRequest(http.MethodGet, desktop.URL(inference.ModelsPrefix), nil)
+			require.NoError(t, err)
+			client.EXPECT().Do(req).Return(test.doResponse, test.doErr)
+
+			originalOsExit := osExit
+			exitCalled := false
+			osExit = func(code int) {
+				exitCalled = true
+				require.Equal(t, 1, code, "Expected exit code to be 1")
+			}
+			defer func() { osExit = originalOsExit }()
+
+			cmd := newStatusCmd(desktop.New(client))
+			buf := new(bytes.Buffer)
+			cmd.SetOut(buf)
+			cmd.SetErr(buf)
+
+			err = cmd.Execute()
+			if test.expectExit {
+				require.True(t, exitCalled, "Expected os.Exit to be called")
+			} else {
+				require.False(t, exitCalled, "Did not expect os.Exit to be called")
+			}
+			if test.expectedErr != nil {
+				require.Error(t, err)
+				require.EqualError(t, err, test.expectedErr.Error())
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, test.expectedOutput, buf.String())
+			}
+		})
+	}
+}