ported workspaceagent test

ethanndickson · ethanndickson · commit b3412661e682 · 2024-06-06T06:12:24.000Z
diff --git a/coderd/workspaceagentsrpc_test.go b/coderd/workspaceagentsrpc_test.go
@@ -1,19 +1,141 @@
 package coderd_test
 
 import (
+	"context"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
+	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/provisionersdk/proto"
 	"github.com/coder/coder/v2/testutil"
 )
 
+// Ported to RPC API from coderd/workspaceagents_test.go
+func TestWorkspaceAgentReportStats(t *testing.T) {
+	t.Parallel()
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		client, db := coderdtest.NewWithDatabase(t, nil)
+		user := coderdtest.CreateFirstUser(t, client)
+		r := dbfake.WorkspaceBuild(t, db, database.Workspace{
+			OrganizationID: user.OrganizationID,
+			OwnerID:        user.UserID,
+		}).WithAgent().Do()
+
+		ac := agentsdk.New(client.URL)
+		ac.SetSessionToken(r.AgentToken)
+		conn, err := ac.ConnectRPC(context.Background())
+		require.NoError(t, err)
+		defer func() {
+			_ = conn.Close()
+		}()
+		agentAPI := agentproto.NewDRPCAgentClient(conn)
+
+		_, err = agentAPI.UpdateStats(context.Background(), &agentproto.UpdateStatsRequest{
+			Stats: &agentproto.Stats{
+				ConnectionsByProto:          map[string]int64{"TCP": 1},
+				ConnectionCount:             1,
+				RxPackets:                   1,
+				RxBytes:                     1,
+				TxPackets:                   1,
+				TxBytes:                     1,
+				SessionCountVscode:          1,
+				SessionCountJetbrains:       0,
+				SessionCountReconnectingPty: 0,
+				SessionCountSsh:             0,
+				ConnectionMedianLatencyMs:   10,
+			},
+		})
+		require.NoError(t, err)
+
+		newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
+		require.NoError(t, err)
+
+		assert.True(t,
+			newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
+			"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
+		)
+	})
+
+	t.Run("FailDeleted", func(t *testing.T) {
+		t.Parallel()
+
+		owner, db := coderdtest.NewWithDatabase(t, nil)
+		ownerUser := coderdtest.CreateFirstUser(t, owner)
+		client, admin := coderdtest.CreateAnotherUser(t, owner, ownerUser.OrganizationID, rbac.RoleTemplateAdmin(), rbac.RoleUserAdmin())
+		r := dbfake.WorkspaceBuild(t, db, database.Workspace{
+			OrganizationID: admin.OrganizationIDs[0],
+			OwnerID:        admin.ID,
+		}).WithAgent().Do()
+
+		ac := agentsdk.New(client.URL)
+		ac.SetSessionToken(r.AgentToken)
+		conn, err := ac.ConnectRPC(context.Background())
+		require.NoError(t, err)
+		defer func() {
+			_ = conn.Close()
+		}()
+		agentAPI := agentproto.NewDRPCAgentClient(conn)
+
+		_, err = agentAPI.UpdateStats(context.Background(), &agentproto.UpdateStatsRequest{
+			Stats: &agentproto.Stats{
+				ConnectionsByProto:          map[string]int64{"TCP": 1},
+				ConnectionCount:             1,
+				RxPackets:                   1,
+				RxBytes:                     1,
+				TxPackets:                   1,
+				TxBytes:                     1,
+				SessionCountVscode:          0,
+				SessionCountJetbrains:       0,
+				SessionCountReconnectingPty: 0,
+				SessionCountSsh:             0,
+				ConnectionMedianLatencyMs:   10,
+			},
+		})
+		require.NoError(t, err)
+
+		newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
+		require.NoError(t, err)
+
+		// nolint:gocritic // using db directly over creating a delete job
+		err = db.UpdateWorkspaceDeletedByID(dbauthz.As(context.Background(),
+			coderdtest.AuthzUserSubject(admin, ownerUser.OrganizationID)),
+			database.UpdateWorkspaceDeletedByIDParams{
+				ID:      newWorkspace.ID,
+				Deleted: true,
+			})
+		require.NoError(t, err)
+
+		resp, err := agentAPI.UpdateStats(context.Background(), &agentproto.UpdateStatsRequest{
+			Stats: &agentproto.Stats{
+				ConnectionsByProto:          map[string]int64{"TCP": 1},
+				ConnectionCount:             1,
+				RxPackets:                   1,
+				RxBytes:                     1,
+				TxPackets:                   1,
+				TxBytes:                     1,
+				SessionCountVscode:          1,
+				SessionCountJetbrains:       0,
+				SessionCountReconnectingPty: 0,
+				SessionCountSsh:             0,
+				ConnectionMedianLatencyMs:   10,
+			},
+		})
+		require.Nil(t, resp)
+		require.ErrorContains(t, err, "agent is invalid")
+	})
+}
+
 func TestAgentAPI_LargeManifest(t *testing.T) {
 	t.Parallel()
 	ctx := testutil.Context(t, testutil.WaitLong)
