@@ -500,7 +500,14 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) {
500500// Verify that the manifest has DisableDirectConnections set to true.
501501agentClient := agentsdk .New (client .URL )
502502agentClient .SetSessionToken (r .AgentToken )
503- manifest := requireGetManifest (ctx ,t ,agentClient )
503+ rpc ,err := agentClient .Listen (ctx )
504+ require .NoError (t ,err )
505+ defer func () {
506+ cErr := rpc .Close ()
507+ require .NoError (t ,cErr )
508+ }()
509+ aAPI := agentproto .NewDRPCAgentClient (rpc )
510+ manifest := requireGetManifest (ctx ,t ,aAPI )
504511require .True (t ,manifest .DisableDirectConnections )
505512
506513_ = agenttest .New (t ,client .URL ,r .AgentToken )
@@ -823,46 +830,63 @@ func TestWorkspaceAgentAppHealth(t *testing.T) {
823830
824831agentClient := agentsdk .New (client .URL )
825832agentClient .SetSessionToken (r .AgentToken )
833+ conn ,err := agentClient .Listen (ctx )
834+ require .NoError (t ,err )
835+ defer func () {
836+ cErr := conn .Close ()
837+ require .NoError (t ,cErr )
838+ }()
839+ aAPI := agentproto .NewDRPCAgentClient (conn )
826840
827- manifest := requireGetManifest (ctx ,t ,agentClient )
841+ manifest := requireGetManifest (ctx ,t ,aAPI )
828842require .EqualValues (t ,codersdk .WorkspaceAppHealthDisabled ,manifest .Apps [0 ].Health )
829843require .EqualValues (t ,codersdk .WorkspaceAppHealthInitializing ,manifest .Apps [1 ].Health )
830- err := agentClient .PostAppHealth (ctx , agentsdk.PostAppHealthsRequest {})
831- require .Error (t ,err )
832844// empty
833- err = agentClient . PostAppHealth (ctx ,agentsdk. PostAppHealthsRequest {})
834- require .Error (t ,err )
845+ _ , err = aAPI . BatchUpdateAppHealths (ctx ,& agentproto. BatchUpdateAppHealthRequest {})
846+ require .NoError (t ,err )
835847// healthcheck disabled
836- err = agentClient .PostAppHealth (ctx , agentsdk.PostAppHealthsRequest {
837- Healths :map [uuid.UUID ]codersdk.WorkspaceAppHealth {
838- manifest .Apps [0 ].ID :codersdk .WorkspaceAppHealthInitializing ,
848+ _ ,err = aAPI .BatchUpdateAppHealths (ctx ,& agentproto.BatchUpdateAppHealthRequest {
849+ Updates : []* agentproto.BatchUpdateAppHealthRequest_HealthUpdate {
850+ {
851+ Id :manifest .Apps [0 ].ID [:],
852+ Health :agentproto .AppHealth_INITIALIZING ,
853+ },
839854},
840855})
841856require .Error (t ,err )
842857// invalid value
843- err = agentClient .PostAppHealth (ctx , agentsdk.PostAppHealthsRequest {
844- Healths :map [uuid.UUID ]codersdk.WorkspaceAppHealth {
845- manifest .Apps [1 ].ID :codersdk .WorkspaceAppHealth ("bad-value" ),
858+ _ ,err = aAPI .BatchUpdateAppHealths (ctx ,& agentproto.BatchUpdateAppHealthRequest {
859+ Updates : []* agentproto.BatchUpdateAppHealthRequest_HealthUpdate {
860+ {
861+ Id :manifest .Apps [1 ].ID [:],
862+ Health :99 ,
863+ },
846864},
847865})
848866require .Error (t ,err )
849867// update to healthy
850- err = agentClient .PostAppHealth (ctx , agentsdk.PostAppHealthsRequest {
851- Healths :map [uuid.UUID ]codersdk.WorkspaceAppHealth {
852- manifest .Apps [1 ].ID :codersdk .WorkspaceAppHealthHealthy ,
868+ _ ,err = aAPI .BatchUpdateAppHealths (ctx ,& agentproto.BatchUpdateAppHealthRequest {
869+ Updates : []* agentproto.BatchUpdateAppHealthRequest_HealthUpdate {
870+ {
871+ Id :manifest .Apps [1 ].ID [:],
872+ Health :agentproto .AppHealth_HEALTHY ,
873+ },
853874},
854875})
855876require .NoError (t ,err )
856- manifest = requireGetManifest (ctx ,t ,agentClient )
877+ manifest = requireGetManifest (ctx ,t ,aAPI )
857878require .EqualValues (t ,codersdk .WorkspaceAppHealthHealthy ,manifest .Apps [1 ].Health )
858879// update to unhealthy
859- err = agentClient .PostAppHealth (ctx , agentsdk.PostAppHealthsRequest {
860- Healths :map [uuid.UUID ]codersdk.WorkspaceAppHealth {
861- manifest .Apps [1 ].ID :codersdk .WorkspaceAppHealthUnhealthy ,
880+ _ ,err = aAPI .BatchUpdateAppHealths (ctx ,& agentproto.BatchUpdateAppHealthRequest {
881+ Updates : []* agentproto.BatchUpdateAppHealthRequest_HealthUpdate {
882+ {
883+ Id :manifest .Apps [1 ].ID [:],
884+ Health :agentproto .AppHealth_UNHEALTHY ,
885+ },
862886},
863887})
864888require .NoError (t ,err )
865- manifest = requireGetManifest (ctx ,t ,agentClient )
889+ manifest = requireGetManifest (ctx ,t ,aAPI )
866890require .EqualValues (t ,codersdk .WorkspaceAppHealthUnhealthy ,manifest .Apps [1 ].Health )
867891}
868892
@@ -1105,8 +1129,15 @@ func TestWorkspaceAgent_Metadata(t *testing.T) {
11051129agentClient .SetSessionToken (r .AgentToken )
11061130
11071131ctx := testutil .Context (t ,testutil .WaitMedium )
1132+ conn ,err := agentClient .Listen (ctx )
1133+ require .NoError (t ,err )
1134+ defer func () {
1135+ cErr := conn .Close ()
1136+ require .NoError (t ,cErr )
1137+ }()
1138+ aAPI := agentproto .NewDRPCAgentClient (conn )
11081139
1109- manifest := requireGetManifest (ctx ,t ,agentClient )
1140+ manifest := requireGetManifest (ctx ,t ,aAPI )
11101141
11111142// Verify manifest API response.
11121143require .Equal (t ,workspace .ID ,manifest .WorkspaceID )
@@ -1276,8 +1307,15 @@ func TestWorkspaceAgent_Metadata_CatchMemoryLeak(t *testing.T) {
12761307agentClient .SetSessionToken (r .AgentToken )
12771308
12781309ctx ,cancel := context .WithCancel (testutil .Context (t ,testutil .WaitSuperLong ))
1310+ conn ,err := agentClient .Listen (ctx )
1311+ require .NoError (t ,err )
1312+ defer func () {
1313+ cErr := conn .Close ()
1314+ require .NoError (t ,cErr )
1315+ }()
1316+ aAPI := agentproto .NewDRPCAgentClient (conn )
12791317
1280- manifest := requireGetManifest (ctx ,t ,agentClient )
1318+ manifest := requireGetManifest (ctx ,t ,aAPI )
12811319
12821320post := func (ctx context.Context ,key ,value string )error {
12831321return agentClient .PostMetadata (ctx , agentsdk.PostMetadataRequest {
@@ -1622,14 +1660,7 @@ func TestWorkspaceAgentExternalAuthListen(t *testing.T) {
16221660})
16231661}
16241662
1625- func requireGetManifest (ctx context.Context ,t testing.TB ,client agent.Client ) agentsdk.Manifest {
1626- conn ,err := client .Listen (ctx )
1627- require .NoError (t ,err )
1628- defer func () {
1629- cErr := conn .Close ()
1630- require .NoError (t ,cErr )
1631- }()
1632- aAPI := agentproto .NewDRPCAgentClient (conn )
1663+ func requireGetManifest (ctx context.Context ,t testing.TB ,aAPI agentproto.DRPCAgentClient ) agentsdk.Manifest {
16331664mp ,err := aAPI .GetManifest (ctx ,& agentproto.GetManifestRequest {})
16341665require .NoError (t ,err )
16351666manifest ,err := agentsdk .ManifestFromProto (mp )