add more tests for get status

Signed-off-by: Gyuho Lee <[email protected]>

client/v1/status.go

@@ -0,0 +1,41 @@
+package v1
+
+import (
+	"context"
+	"crypto/tls"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+
+	"github.com/leptonai/gpud/manager/packages"
+)
+
+func GetStatus(ctx context.Context, url string) ([]packages.PackageStatus, error) {
+	httpClient := &http.Client{
+		Transport: &http.Transport{
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		},
+	}
+	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := httpClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("unexpected status code %v received", resp.StatusCode)
+	}
+	rawBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+	var ret []packages.PackageStatus
+	if err := json.Unmarshal(rawBody, &ret); err != nil {
+		return nil, err
+	}
+	return ret, nil
+}

client/v1/status_test.go

@@ -0,0 +1,115 @@
+package v1
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/leptonai/gpud/manager/packages"
+)
+
+func TestGetStatus(t *testing.T) {
+	tests := []struct {
+		name           string
+		serverResponse func(w http.ResponseWriter, r *http.Request)
+		wantErr        bool
+		expectedData   []packages.PackageStatus
+	}{
+		{
+			name: "successful response",
+			serverResponse: func(w http.ResponseWriter, r *http.Request) {
+				assert.Equal(t, http.MethodGet, r.Method)
+
+				response := []packages.PackageStatus{
+					{
+						Name:           "test-package",
+						IsInstalled:    true,
+						Installing:     false,
+						Progress:       100,
+						TotalTime:      1 * time.Hour,
+						Status:         true,
+						TargetVersion:  "1.0.0",
+						CurrentVersion: "1.0.0",
+						ScriptPath:     "/path/to/script",
+						Dependency:     [][]string{{"dep1", "1.0.0"}},
+					},
+				}
+				json.NewEncoder(w).Encode(response)
+			},
+			wantErr: false,
+			expectedData: []packages.PackageStatus{
+				{
+					Name:           "test-package",
+					IsInstalled:    true,
+					Installing:     false,
+					Progress:       100,
+					TotalTime:      1 * time.Hour,
+					Status:         true,
+					TargetVersion:  "1.0.0",
+					CurrentVersion: "1.0.0",
+					ScriptPath:     "/path/to/script",
+					Dependency:     [][]string{{"dep1", "1.0.0"}},
+				},
+			},
+		},
+		{
+			name: "server returns error status",
+			serverResponse: func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusInternalServerError)
+			},
+			wantErr:      true,
+			expectedData: nil,
+		},
+		{
+			name: "invalid JSON response",
+			serverResponse: func(w http.ResponseWriter, r *http.Request) {
+				w.Write([]byte("invalid json"))
+			},
+			wantErr:      true,
+			expectedData: nil,
+		},
+		{
+			name: "context canceled",
+			serverResponse: func(w http.ResponseWriter, r *http.Request) {
+				// Simulate a slow response that will be canceled
+				time.Sleep(100 * time.Millisecond)
+				json.NewEncoder(w).Encode([]packages.PackageStatus{})
+			},
+			wantErr:      true,
+			expectedData: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create test server
+			server := httptest.NewTLSServer(http.HandlerFunc(tt.serverResponse))
+			defer server.Close()
+
+			var ctx context.Context
+			var cancel context.CancelFunc
+			if tt.name == "context canceled" {
+				ctx, cancel = context.WithTimeout(context.Background(), 50*time.Millisecond)
+			} else {
+				ctx, cancel = context.WithTimeout(context.Background(), time.Second)
+			}
+			defer cancel()
+
+			// Call GetStatus with the test server's URL
+			status, err := GetStatus(ctx, server.URL)
+
+			if tt.wantErr {
+				assert.Error(t, err)
+				assert.Nil(t, status)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.expectedData, status)
+			}
+		})
+	}
+}

cmd/gpud/command/status.go

@@ -2,22 +2,17 @@ package command
 
 import (
 	"context"
-	"crypto/tls"
-	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
-	"net/http"
 	"time"
 
+	"github.com/urfave/cli"
+
 	client "github.com/leptonai/gpud/client/v1"
 	"github.com/leptonai/gpud/config"
 	"github.com/leptonai/gpud/errdefs"
 	"github.com/leptonai/gpud/log"
-	"github.com/leptonai/gpud/manager/packages"
 	"github.com/leptonai/gpud/pkg/systemd"
-
-	"github.com/urfave/cli"
 )
 
 func cmdStatus(cliContext *cli.Context) error {
@@ -56,7 +51,9 @@ func cmdStatus(cliContext *cli.Context) error {
 	fmt.Printf("%s successfully checked gpud health\n", checkMark)
 
 	for {
-		packageStatus, err := getStatus()
+		cctx, ccancel := context.WithTimeout(rootCtx, 15*time.Second)
+		packageStatus, err := client.GetStatus(cctx, fmt.Sprintf("https://localhost:%d/admin/packages", config.DefaultGPUdPort))
+		ccancel()
 		if err != nil {
 			fmt.Printf("%s failed to get package status: %v\n", warningSign, err)
 			return err
@@ -141,32 +138,3 @@ func checkNvidiaInfoComponent() error {
 
 	return nil
 }
-
-func getStatus() ([]packages.PackageStatus, error) {
-	httpClient := &http.Client{
-		Transport: &http.Transport{
-			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
-		},
-	}
-	req, err := http.NewRequest("GET", fmt.Sprintf("https://localhost:%d/admin/packages", config.DefaultGPUdPort), nil)
-	if err != nil {
-		return nil, err
-	}
-	resp, err := httpClient.Do(req)
-	if err != nil {
-		return nil, err
-	}
-	defer resp.Body.Close()
-	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("unexpected status code %v received", resp.StatusCode)
-	}
-	rawBody, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
-	}
-	var ret []packages.PackageStatus
-	if err := json.Unmarshal(rawBody, &ret); err != nil {
-		return nil, err
-	}
-	return ret, nil
-}