fix(internal/server): return 404 not found on unknown component queries

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

client/v1/v1.go

@@ -11,6 +11,7 @@ import (
 	"net/url"
 
 	v1 "github.com/leptonai/gpud/api/v1"
+	"github.com/leptonai/gpud/errdefs"
 	"github.com/leptonai/gpud/internal/server"
 	"sigs.k8s.io/yaml"
 )
@@ -213,7 +214,11 @@ func GetStates(ctx context.Context, addr string, opts ...OpOption) (v1.LeptonSta
 		return nil, fmt.Errorf("failed to make request: %w", err)
 	}
 	defer resp.Body.Close()
+
 	if resp.StatusCode != http.StatusOK {
+		if resp.StatusCode == http.StatusNotFound {
+			return nil, errdefs.ErrNotFound
+		}
 		return nil, errors.New("server not ready, response not 200")
 	}
 

components/components.go

@@ -127,13 +127,19 @@ func GetComponent(name string) (Component, error) {
 	defaultSetMu.RLock()
 	defer defaultSetMu.RUnlock()
 
-	if defaultSet == nil {
+	return getComponent(defaultSet, name)
+}
+
+func getComponent(set map[string]Component, name string) (Component, error) {
+	if set == nil {
 		return nil, fmt.Errorf("component set not initialized: %w", errdefs.ErrUnavailable)
 	}
-	if _, ok := defaultSet[name]; !ok {
+
+	v, ok := set[name]
+	if !ok {
 		return nil, fmt.Errorf("component %s not found: %w", name, errdefs.ErrNotFound)
 	}
-	return defaultSet[name], nil
+	return v, nil
 }
 
 func GetAllComponents() map[string]Component {

components/components_test.go

@@ -0,0 +1,17 @@
+package components
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/leptonai/gpud/errdefs"
+)
+
+func TestGetComponentErrors(t *testing.T) {
+	if _, err := getComponent(nil, "nvidia"); !errors.Is(err, errdefs.ErrUnavailable) {
+		t.Errorf("expected ErrUnavailable, got %v", err)
+	}
+	if _, err := getComponent(map[string]Component{}, "nvidia"); !errors.Is(err, errdefs.ErrNotFound) {
+		t.Errorf("expected ErrNotFound, got %v", err)
+	}
+}

e2e/e2e_test.go

@@ -7,6 +7,7 @@ import (
 	"crypto/tls"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -18,6 +19,7 @@ import (
 
 	v1 "github.com/leptonai/gpud/api/v1"
 	client_v1 "github.com/leptonai/gpud/client/v1"
+	"github.com/leptonai/gpud/errdefs"
 	"github.com/leptonai/gpud/internal/server"
 )
 
@@ -314,6 +316,10 @@ func TestGpudHealthzInfo(t *testing.T) {
 			}
 		}
 
+		if _, err := client_v1.GetStates(ctx, "https://"+ep, append(opts, client_v1.WithComponent("unknown!!!"))...); !errors.Is(err, errdefs.ErrNotFound) {
+			t.Errorf("expected ErrNotFound, got %v", err)
+		}
+
 		states, err := client_v1.GetStates(ctx, "https://"+ep, opts...)
 		if err != nil {
 			t.Errorf("failed to get states: %v", err)

internal/server/handlers.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
 	"sort"
@@ -72,7 +73,7 @@ func (g *globalHandler) getReqComponents(c *gin.Context) ([]string, error) {
 	var ret []string
 	for _, component := range strings.Split(components, ",") {
 		if _, err := lep_components.GetComponent(component); err != nil {
-			return nil, fmt.Errorf("failed to get component: %v", err)
+			return nil, fmt.Errorf("failed to get component: %v (%w)", err, errors.Unwrap(err))
 		}
 		ret = append(ret, component)
 	}

internal/server/handlers_components.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"errors"
 	"net/http"
 	"sort"
 	"time"
@@ -122,6 +123,11 @@ func (g *globalHandler) getStates(c *gin.Context) {
 	var states v1.LeptonStates
 	components, err := g.getReqComponents(c)
 	if err != nil {
+		if errors.Is(err, errdefs.ErrNotFound) {
+			c.JSON(http.StatusNotFound, gin.H{"code": errdefs.ErrNotFound, "message": "component not found: " + err.Error()})
+			return
+		}
+
 		c.JSON(http.StatusBadRequest, gin.H{"code": errdefs.ErrInvalidArgument, "message": "failed to parse components: " + err.Error()})
 		return
 	}
@@ -193,6 +199,11 @@ func (g *globalHandler) getEvents(c *gin.Context) {
 	var events v1.LeptonEvents
 	components, err := g.getReqComponents(c)
 	if err != nil {
+		if errors.Is(err, errdefs.ErrNotFound) {
+			c.JSON(http.StatusNotFound, gin.H{"code": errdefs.ErrNotFound, "message": "component not found: " + err.Error()})
+			return
+		}
+
 		c.JSON(http.StatusBadRequest, gin.H{"code": errdefs.ErrInvalidArgument, "message": "failed to parse components: " + err.Error()})
 		return
 	}
@@ -270,6 +281,11 @@ func (g *globalHandler) getInfo(c *gin.Context) {
 	var infos v1.LeptonInfo
 	components, err := g.getReqComponents(c)
 	if err != nil {
+		if errors.Is(err, errdefs.ErrNotFound) {
+			c.JSON(http.StatusNotFound, gin.H{"code": errdefs.ErrNotFound, "message": "component not found: " + err.Error()})
+			return
+		}
+
 		c.JSON(http.StatusBadRequest, gin.H{"code": errdefs.ErrInvalidArgument, "message": "failed to parse components: " + err.Error()})
 		return
 	}
@@ -377,6 +393,11 @@ const (
 func (g *globalHandler) getMetrics(c *gin.Context) {
 	components, err := g.getReqComponents(c)
 	if err != nil {
+		if errors.Is(err, errdefs.ErrNotFound) {
+			c.JSON(http.StatusNotFound, gin.H{"code": errdefs.ErrNotFound, "message": "component not found: " + err.Error()})
+			return
+		}
+
 		c.JSON(http.StatusBadRequest, gin.H{"code": errdefs.ErrInvalidArgument, "message": "failed to parse components: " + err.Error()})
 		return
 	}