### Question:
QueryContext { public synchronized void setMemoryPool(MemoryPool newMemoryPool) { requireNonNull(newMemoryPool, "newMemoryPool is null"); if (memoryPool == newMemoryPool) { return; } ListenableFuture<?> future = memoryPool.moveQuery(queryId, newMemoryPool); memoryPool = newMemoryPool; future.addListener(() -> { taskContexts.values().forEach(TaskContext::moreMemoryAvailable); }, directExecutor()); } QueryContext(
            QueryId queryId,
            DataSize maxUserMemory,
            DataSize maxTotalMemory,
            MemoryPool memoryPool,
            GcMonitor gcMonitor,
            Executor notificationExecutor,
            ScheduledExecutorService yieldExecutor,
            DataSize maxSpill,
            SpillSpaceTracker spillSpaceTracker); synchronized void setResourceOvercommit(); synchronized ListenableFuture<?> reserveSpill(long bytes); synchronized void freeSpill(long bytes); synchronized void setMemoryPool(MemoryPool newMemoryPool); synchronized void setMaxUserMemory(long maxUserMemory); synchronized void setMaxTotalMemory(long maxTotalMemory); synchronized MemoryPool getMemoryPool(); TaskContext addTaskContext(
            TaskStateMachine taskStateMachine,
            Session session,
            Runnable notifyStatusChanged,
            boolean perOperatorCpuTimerEnabled,
            boolean cpuTimerEnabled,
            OptionalInt totalPartitions); R accept(QueryContextVisitor<C, R> visitor, C context); List<R> acceptChildren(QueryContextVisitor<C, R> visitor, C context); TaskContext getTaskContextByTaskId(TaskId taskId);  }### Answer:
@Test(dataProvider = "testSetMemoryPoolOptions") public void testSetMemoryPool(boolean useReservedPool) { QueryId secondQuery = new QueryId("second"); MemoryPool reservedPool = new MemoryPool(RESERVED_POOL, DataSize.ofBytes(10)); long secondQueryMemory = reservedPool.getMaxBytes() - 1; if (useReservedPool) { assertTrue(reservedPool.reserve(secondQuery, "test", secondQueryMemory).isDone()); } try (LocalQueryRunner localQueryRunner = LocalQueryRunner.create(TEST_SESSION)) { QueryContext queryContext = new QueryContext( new QueryId("query"), DataSize.ofBytes(10), DataSize.ofBytes(20), new MemoryPool(GENERAL_POOL, DataSize.ofBytes(10)), new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), DataSize.ofBytes(0), new SpillSpaceTracker(DataSize.ofBytes(0))); queryContext.getQueryMemoryContext().initializeLocalMemoryContexts("test"); LocalMemoryContext userMemoryContext = queryContext.getQueryMemoryContext().localUserMemoryContext(); LocalMemoryContext revocableMemoryContext = queryContext.getQueryMemoryContext().localRevocableMemoryContext(); assertTrue(userMemoryContext.setBytes(3).isDone()); assertTrue(revocableMemoryContext.setBytes(5).isDone()); queryContext.setMemoryPool(reservedPool); if (useReservedPool) { reservedPool.free(secondQuery, "test", secondQueryMemory); } userMemoryContext.close(); revocableMemoryContext.close(); } }