mirror of https://github.com/coder/coder.git
277 lines
7.4 KiB
Go
277 lines
7.4 KiB
Go
package agentapi_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/mock/gomock"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
|
|
"cdr.dev/slog/sloggers/slogtest"
|
|
|
|
agentproto "github.com/coder/coder/v2/agent/proto"
|
|
"github.com/coder/coder/v2/coderd/agentapi"
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/database/dbmock"
|
|
"github.com/coder/coder/v2/coderd/database/dbtime"
|
|
"github.com/coder/coder/v2/coderd/database/pubsub"
|
|
)
|
|
|
|
type fakePublisher struct {
|
|
// Nil pointer to pass interface check.
|
|
pubsub.Pubsub
|
|
publishes [][]byte
|
|
}
|
|
|
|
var _ pubsub.Pubsub = &fakePublisher{}
|
|
|
|
func (f *fakePublisher) Publish(_ string, message []byte) error {
|
|
f.publishes = append(f.publishes, message)
|
|
return nil
|
|
}
|
|
|
|
func TestBatchUpdateMetadata(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
agent := database.WorkspaceAgent{
|
|
ID: uuid.New(),
|
|
}
|
|
|
|
t.Run("OK", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
dbM := dbmock.NewMockStore(gomock.NewController(t))
|
|
pub := &fakePublisher{}
|
|
|
|
now := dbtime.Now()
|
|
req := &agentproto.BatchUpdateMetadataRequest{
|
|
Metadata: []*agentproto.Metadata{
|
|
{
|
|
Key: "awesome key",
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
CollectedAt: timestamppb.New(now.Add(-10 * time.Second)),
|
|
Age: 10,
|
|
Value: "awesome value",
|
|
Error: "",
|
|
},
|
|
},
|
|
{
|
|
Key: "uncool key",
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
CollectedAt: timestamppb.New(now.Add(-3 * time.Second)),
|
|
Age: 3,
|
|
Value: "",
|
|
Error: "uncool value",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
dbM.EXPECT().UpdateWorkspaceAgentMetadata(gomock.Any(), database.UpdateWorkspaceAgentMetadataParams{
|
|
WorkspaceAgentID: agent.ID,
|
|
Key: []string{req.Metadata[0].Key, req.Metadata[1].Key},
|
|
Value: []string{req.Metadata[0].Result.Value, req.Metadata[1].Result.Value},
|
|
Error: []string{req.Metadata[0].Result.Error, req.Metadata[1].Result.Error},
|
|
// The value from the agent is ignored.
|
|
CollectedAt: []time.Time{now, now},
|
|
}).Return(nil)
|
|
|
|
api := &agentapi.MetadataAPI{
|
|
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
|
|
return agent, nil
|
|
},
|
|
Database: dbM,
|
|
Pubsub: pub,
|
|
Log: slogtest.Make(t, nil),
|
|
TimeNowFn: func() time.Time {
|
|
return now
|
|
},
|
|
}
|
|
|
|
resp, err := api.BatchUpdateMetadata(context.Background(), req)
|
|
require.NoError(t, err)
|
|
require.Equal(t, &agentproto.BatchUpdateMetadataResponse{}, resp)
|
|
|
|
require.Equal(t, 1, len(pub.publishes))
|
|
var gotEvent agentapi.WorkspaceAgentMetadataChannelPayload
|
|
require.NoError(t, json.Unmarshal(pub.publishes[0], &gotEvent))
|
|
require.Equal(t, agentapi.WorkspaceAgentMetadataChannelPayload{
|
|
CollectedAt: now,
|
|
Keys: []string{req.Metadata[0].Key, req.Metadata[1].Key},
|
|
}, gotEvent)
|
|
})
|
|
|
|
t.Run("ExceededLength", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
dbM := dbmock.NewMockStore(gomock.NewController(t))
|
|
pub := pubsub.NewInMemory()
|
|
|
|
almostLongValue := ""
|
|
for i := 0; i < 2048; i++ {
|
|
almostLongValue += "a"
|
|
}
|
|
|
|
now := dbtime.Now()
|
|
req := &agentproto.BatchUpdateMetadataRequest{
|
|
Metadata: []*agentproto.Metadata{
|
|
{
|
|
Key: "almost long value",
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
Value: almostLongValue,
|
|
},
|
|
},
|
|
{
|
|
Key: "too long value",
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
Value: almostLongValue + "a",
|
|
},
|
|
},
|
|
{
|
|
Key: "almost long error",
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
Error: almostLongValue,
|
|
},
|
|
},
|
|
{
|
|
Key: "too long error",
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
Error: almostLongValue + "a",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
dbM.EXPECT().UpdateWorkspaceAgentMetadata(gomock.Any(), database.UpdateWorkspaceAgentMetadataParams{
|
|
WorkspaceAgentID: agent.ID,
|
|
Key: []string{req.Metadata[0].Key, req.Metadata[1].Key, req.Metadata[2].Key, req.Metadata[3].Key},
|
|
Value: []string{
|
|
almostLongValue,
|
|
almostLongValue, // truncated
|
|
"",
|
|
"",
|
|
},
|
|
Error: []string{
|
|
"",
|
|
"value of 2049 bytes exceeded 2048 bytes",
|
|
almostLongValue,
|
|
"error of 2049 bytes exceeded 2048 bytes", // replaced
|
|
},
|
|
// The value from the agent is ignored.
|
|
CollectedAt: []time.Time{now, now, now, now},
|
|
}).Return(nil)
|
|
|
|
api := &agentapi.MetadataAPI{
|
|
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
|
|
return agent, nil
|
|
},
|
|
Database: dbM,
|
|
Pubsub: pub,
|
|
Log: slogtest.Make(t, nil),
|
|
TimeNowFn: func() time.Time {
|
|
return now
|
|
},
|
|
}
|
|
|
|
resp, err := api.BatchUpdateMetadata(context.Background(), req)
|
|
require.NoError(t, err)
|
|
require.Equal(t, &agentproto.BatchUpdateMetadataResponse{}, resp)
|
|
})
|
|
|
|
t.Run("KeysTooLong", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
dbM := dbmock.NewMockStore(gomock.NewController(t))
|
|
pub := pubsub.NewInMemory()
|
|
|
|
now := dbtime.Now()
|
|
req := &agentproto.BatchUpdateMetadataRequest{
|
|
Metadata: []*agentproto.Metadata{
|
|
{
|
|
Key: "key 1",
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
Value: "value 1",
|
|
},
|
|
},
|
|
{
|
|
Key: "key 2",
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
Value: "value 2",
|
|
},
|
|
},
|
|
{
|
|
Key: func() string {
|
|
key := "key 3 "
|
|
for i := 0; i < (6144 - len("key 1") - len("key 2") - len("key 3") - 1); i++ {
|
|
key += "a"
|
|
}
|
|
return key
|
|
}(),
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
Value: "value 3",
|
|
},
|
|
},
|
|
{
|
|
Key: "a", // should be ignored
|
|
Result: &agentproto.WorkspaceAgentMetadata_Result{
|
|
Value: "value 4",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
dbM.EXPECT().UpdateWorkspaceAgentMetadata(gomock.Any(), database.UpdateWorkspaceAgentMetadataParams{
|
|
WorkspaceAgentID: agent.ID,
|
|
// No key 4.
|
|
Key: []string{req.Metadata[0].Key, req.Metadata[1].Key, req.Metadata[2].Key},
|
|
Value: []string{req.Metadata[0].Result.Value, req.Metadata[1].Result.Value, req.Metadata[2].Result.Value},
|
|
Error: []string{req.Metadata[0].Result.Error, req.Metadata[1].Result.Error, req.Metadata[2].Result.Error},
|
|
// The value from the agent is ignored.
|
|
CollectedAt: []time.Time{now, now, now},
|
|
}).Return(nil)
|
|
|
|
api := &agentapi.MetadataAPI{
|
|
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
|
|
return agent, nil
|
|
},
|
|
Database: dbM,
|
|
Pubsub: pub,
|
|
Log: slogtest.Make(t, nil),
|
|
TimeNowFn: func() time.Time {
|
|
return now
|
|
},
|
|
}
|
|
|
|
// Watch the pubsub for events.
|
|
var (
|
|
eventCount int64
|
|
gotEvent agentapi.WorkspaceAgentMetadataChannelPayload
|
|
)
|
|
cancel, err := pub.Subscribe(agentapi.WatchWorkspaceAgentMetadataChannel(agent.ID), func(ctx context.Context, message []byte) {
|
|
if atomic.AddInt64(&eventCount, 1) > 1 {
|
|
return
|
|
}
|
|
require.NoError(t, json.Unmarshal(message, &gotEvent))
|
|
})
|
|
require.NoError(t, err)
|
|
defer cancel()
|
|
|
|
resp, err := api.BatchUpdateMetadata(context.Background(), req)
|
|
require.Error(t, err)
|
|
require.Equal(t, "metadata keys of 6145 bytes exceeded 6144 bytes", err.Error())
|
|
require.Nil(t, resp)
|
|
|
|
require.Equal(t, int64(1), atomic.LoadInt64(&eventCount))
|
|
require.Equal(t, agentapi.WorkspaceAgentMetadataChannelPayload{
|
|
CollectedAt: now,
|
|
// No key 4.
|
|
Keys: []string{req.Metadata[0].Key, req.Metadata[1].Key, req.Metadata[2].Key},
|
|
}, gotEvent)
|
|
})
|
|
}
|