feat: agentapi uses appearance.Fetcher

spikecurtis · spikecurtis · commit 385e74a9405e · 2024-01-29T14:43:51.000+04:00
diff --git a/agent/proto/convert.go b/agent/proto/convert.go
@@ -104,3 +104,19 @@ func SDKAppFromProto(protoApp *WorkspaceApp) (codersdk.WorkspaceApp, error) {
 		Health: health,
 	}, nil
 }
+
+func SDKServiceBannerFromProto(sbp *ServiceBanner) codersdk.ServiceBannerConfig {
+	return codersdk.ServiceBannerConfig{
+		Enabled:         sbp.GetEnabled(),
+		Message:         sbp.GetMessage(),
+		BackgroundColor: sbp.GetBackgroundColor(),
+	}
+}
+
+func ServiceBannerFromSDK(sb codersdk.ServiceBannerConfig) *ServiceBanner {
+	return &ServiceBanner{
+		Enabled:         sb.Enabled,
+		Message:         sb.Message,
+		BackgroundColor: sb.BackgroundColor,
+	}
+}
diff --git a/coderd/agentapi/api.go b/coderd/agentapi/api.go
@@ -17,6 +17,7 @@ import (
 
 	"cdr.dev/slog"
 	agentproto "github.com/coder/coder/v2/agent/proto"
+	"github.com/coder/coder/v2/coderd/appearance"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/externalauth"
@@ -61,6 +62,7 @@ type Options struct {
 	TailnetCoordinator                *atomic.Pointer[tailnet.Coordinator]
 	TemplateScheduleStore             *atomic.Pointer[schedule.TemplateScheduleStore]
 	StatsBatcher                      StatsBatcher
+	AppearanceFetcher                 *atomic.Pointer[appearance.Fetcher]
 	PublishWorkspaceUpdateFn          func(ctx context.Context, workspaceID uuid.UUID)
 	PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
 
@@ -98,7 +100,7 @@ func New(opts Options) *API {
 	}
 
 	api.ServiceBannerAPI = &ServiceBannerAPI{
-		Database: opts.Database,
+		appearanceFetcher: opts.AppearanceFetcher,
 	}
 
 	api.StatsAPI = &StatsAPI{
diff --git a/coderd/agentapi/servicebanner.go b/coderd/agentapi/servicebanner.go
@@ -2,37 +2,22 @@ package agentapi
 
 import (
 	"context"
-	"database/sql"
-	"encoding/json"
+	"sync/atomic"
 
 	"golang.org/x/xerrors"
 
-	agentproto "github.com/coder/coder/v2/agent/proto"
-	"github.com/coder/coder/v2/coderd/database"
-	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/agent/proto"
+	"github.com/coder/coder/v2/coderd/appearance"
 )
 
 type ServiceBannerAPI struct {
-	Database database.Store
+	appearanceFetcher *atomic.Pointer[appearance.Fetcher]
 }
 
-func (a *ServiceBannerAPI) GetServiceBanner(ctx context.Context, _ *agentproto.GetServiceBannerRequest) (*agentproto.ServiceBanner, error) {
-	serviceBannerJSON, err := a.Database.GetServiceBanner(ctx)
-	if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
-		return nil, xerrors.Errorf("get service banner: %w", err)
+func (a *ServiceBannerAPI) GetServiceBanner(ctx context.Context, _ *proto.GetServiceBannerRequest) (*proto.ServiceBanner, error) {
+	cfg, err := (*a.appearanceFetcher.Load()).Fetch(ctx)
+	if err != nil {
+		return nil, xerrors.Errorf("fetch appearance: %w", err)
 	}
-
-	var cfg codersdk.ServiceBannerConfig
-	if serviceBannerJSON != "" {
-		err = json.Unmarshal([]byte(serviceBannerJSON), &cfg)
-		if err != nil {
-			return nil, xerrors.Errorf("unmarshal json: %w, raw: %s", err, serviceBannerJSON)
-		}
-	}
-
-	return &agentproto.ServiceBanner{
-		Enabled:         cfg.Enabled,
-		Message:         cfg.Message,
-		BackgroundColor: cfg.BackgroundColor,
-	}, nil
+	return proto.ServiceBannerFromSDK(cfg.ServiceBanner), nil
 }
diff --git a/coderd/agentapi/servicebanner_internal_test.go b/coderd/agentapi/servicebanner_internal_test.go
@@ -0,0 +1,72 @@
+package agentapi
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+
+	"golang.org/x/xerrors"
+
+	agentproto "github.com/coder/coder/v2/agent/proto"
+	"github.com/coder/coder/v2/coderd/appearance"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetServiceBanner(t *testing.T) {
+	t.Parallel()
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		cfg := codersdk.ServiceBannerConfig{
+			Enabled:         true,
+			Message:         "hello world",
+			BackgroundColor: "#000000",
+		}
+
+		var ff appearance.Fetcher = fakeFetcher{cfg: codersdk.AppearanceConfig{ServiceBanner: cfg}}
+		ptr := atomic.Pointer[appearance.Fetcher]{}
+		ptr.Store(&ff)
+
+		api := &ServiceBannerAPI{
+			appearanceFetcher: &ptr,
+		}
+
+		resp, err := api.GetServiceBanner(context.Background(), &agentproto.GetServiceBannerRequest{})
+		require.NoError(t, err)
+
+		require.Equal(t, &agentproto.ServiceBanner{
+			Enabled:         cfg.Enabled,
+			Message:         cfg.Message,
+			BackgroundColor: cfg.BackgroundColor,
+		}, resp)
+	})
+
+	t.Run("FetchError", func(t *testing.T) {
+		t.Parallel()
+
+		expectedErr := xerrors.New("badness")
+		var ff appearance.Fetcher = fakeFetcher{err: expectedErr}
+		ptr := atomic.Pointer[appearance.Fetcher]{}
+		ptr.Store(&ff)
+
+		api := &ServiceBannerAPI{
+			appearanceFetcher: &ptr,
+		}
+
+		resp, err := api.GetServiceBanner(context.Background(), &agentproto.GetServiceBannerRequest{})
+		require.Error(t, err)
+		require.ErrorIs(t, err, expectedErr)
+		require.Nil(t, resp)
+	})
+}
+
+type fakeFetcher struct {
+	cfg codersdk.AppearanceConfig
+	err error
+}
+
+func (f fakeFetcher) Fetch(context.Context) (codersdk.AppearanceConfig, error) {
+	return f.cfg, f.err
+}
diff --git a/coderd/agentapi/servicebanner_test.go b/coderd/agentapi/servicebanner_test.go
diff --git a/coderd/workspaceagentsrpc.go b/coderd/workspaceagentsrpc.go
@@ -130,6 +130,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 		DerpMapFn:                         api.DERPMap,
 		TailnetCoordinator:                &api.TailnetCoordinator,
 		TemplateScheduleStore:             api.TemplateScheduleStore,
+		AppearanceFetcher:                 &api.AppearanceFetcher,
 		StatsBatcher:                      api.statsBatcher,
 		PublishWorkspaceUpdateFn:          api.publishWorkspaceUpdate,
 		PublishWorkspaceAgentLogsUpdateFn: api.publishWorkspaceAgentLogsUpdate,
diff --git a/enterprise/coderd/appearance_test.go b/enterprise/coderd/appearance_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/cli/clibase"
 	"github.com/coder/coder/v2/coderd/appearance"
 	"github.com/coder/coder/v2/coderd/coderdtest"
@@ -159,6 +160,8 @@ func TestServiceBanners(t *testing.T) {
 		banner, err := agentClient.GetServiceBanner(ctx)
 		require.NoError(t, err)
 		require.Equal(t, cfg.ServiceBanner, banner)
+		banner = requireGetServiceBannerV2(ctx, t, agentClient)
+		require.Equal(t, cfg.ServiceBanner, banner)
 
 		// Create an AGPL Coderd against the same database
 		agplClient := coderdtest.New(t, &coderdtest.Options{Database: store, Pubsub: ps})
@@ -167,16 +170,32 @@ func TestServiceBanners(t *testing.T) {
 		banner, err = agplAgentClient.GetServiceBanner(ctx)
 		require.NoError(t, err)
 		require.Equal(t, codersdk.ServiceBannerConfig{}, banner)
+		banner = requireGetServiceBannerV2(ctx, t, agplAgentClient)
+		require.Equal(t, codersdk.ServiceBannerConfig{}, banner)
 
 		// No license means no banner.
 		err = client.DeleteLicense(ctx, lic.ID)
 		require.NoError(t, err)
 		banner, err = agentClient.GetServiceBanner(ctx)
 		require.NoError(t, err)
 		require.Equal(t, codersdk.ServiceBannerConfig{}, banner)
+		banner = requireGetServiceBannerV2(ctx, t, agentClient)
+		require.Equal(t, codersdk.ServiceBannerConfig{}, banner)
 	})
 }
 
+func requireGetServiceBannerV2(ctx context.Context, t *testing.T, client *agentsdk.Client) codersdk.ServiceBannerConfig {
+	cc, err := client.Listen(ctx)
+	require.NoError(t, err)
+	defer func() {
+		_ = cc.Close()
+	}()
+	aAPI := proto.NewDRPCAgentClient(cc)
+	sbp, err := aAPI.GetServiceBanner(ctx, &proto.GetServiceBannerRequest{})
+	require.NoError(t, err)
+	return proto.SDKServiceBannerFromProto(sbp)
+}
+
 func TestCustomSupportLinks(t *testing.T) {
 	t.Parallel()
 
