@Override
    public void registerQuery(ManagedQueryExecution queryExecution)
    {
        QueryId queryId = queryExecution.getBasicQueryInfo().getQueryId();
        queries.computeIfAbsent(queryId, unused -> {
            AtomicLong sequenceId = new AtomicLong();
            PeriodicTaskExecutor taskExecutor = new PeriodicTaskExecutor(
                    queryHeartbeatInterval.toMillis(),
                    executor,
                    () -> sendQueryHeartbeat(queryExecution, sequenceId.incrementAndGet()));
            taskExecutor.start();
            return taskExecutor;
        });
        queryExecution.addStateChangeListener(newState -> {
            if (newState.isDone()) {
                queries.computeIfPresent(queryId, (unused, queryHeartbeatSender) -> {
                    queryHeartbeatSender.forceRun();
                    queryHeartbeatSender.stop();
                    return null;
                });
            }
        });
    }
@Test(timeOut = 6_000)
    public void testQueryHeartbeat()
            throws Exception
    {
        MockManagedQueryExecution queryExecution = new MockManagedQueryExecution(1);
        sender.registerQuery(queryExecution);

        Thread.sleep(SLEEP_DURATION);

        int queryHeartbeats = resourceManagerClient.getQueryHeartbeats();
        assertTrue(queryHeartbeats > TARGET_HEARTBEATS * 0.5 && queryHeartbeats <= TARGET_HEARTBEATS * 1.5,
                format("Expect number of heartbeats to fall within target range (%s), +/- 50%%.  Was: %s", TARGET_HEARTBEATS, queryHeartbeats));

        // Completing the query stops the heartbeats
        queryExecution.complete();
        queryHeartbeats = resourceManagerClient.getQueryHeartbeats();

        Thread.sleep(SLEEP_DURATION);

        assertTrue(resourceManagerClient.getQueryHeartbeats() <= (queryHeartbeats + 1));
    }