fix test and dispatcher for pings

Overlord-Client/cmd/agent/handlers/dispatcher_test.go

@@ -45,8 +45,14 @@ func TestHandlePing(t *testing.T) {
 		t.Fatalf("HandlePing failed: %v", err)
 	}
 
-	if len(writer.msgs) < 1 {
-		t.Fatal("Expected at least 1 message")
+	deadline := time.After(2 * time.Second)
+	for len(writer.msgs) < 1 {
+		select {
+		case <-deadline:
+			t.Fatal("Timed out waiting for pong message")
+		default:
+			time.Sleep(5 * time.Millisecond)
+		}
 	}
 
 	var pong wire.Pong

Overlord-Client/cmd/agent/handlers/ping.go

@@ -14,6 +14,8 @@ func HandlePing(ctx context.Context, env *runtime.Env, envelope map[string]inter
 	if ts == 0 {
 		ts = time.Now().UnixMilli()
 	}
+	env.SetLastPong(time.Now().UnixMilli())
+
 	pong := wire.Pong{Type: "pong", TS: ts}
 	go func() {
 		if err := wire.WriteMsg(ctx, env.Conn, pong); err != nil {