private Future<?> drain(String workerId) {
        return scheduleInternal(() -> {
            workerStatsManager.drainTotalExecTimeStart();
            assignmentsMovedInLastDrain = invokeDrain(workerId);
            workerStatsManager.drainTotalExecTimeEnd();
        }, "Encountered error when invoking drain");
    }
@Test
    public void testDrain() throws Exception {

        final int numWorkers = 4;

        List<WorkerInfo> workerInfoList = new LinkedList<>();
        for (int ix = 0; ix < numWorkers; ix++) {
            String workerId = "worker-" + ix;
            String workerHostName = "workerHostname-" + ix;
            workerInfoList.add(WorkerInfo.of(workerId, workerHostName, 5000));
        }
        doReturn(workerInfoList).when(membershipManager).getCurrentMembership();

        // Set up multiple functions and assignments, so that there are more functions than workers.
        final int numFunctions = numWorkers * 5;
        List<Function.FunctionMetaData> functionMetaDataList = new LinkedList<>();
        long version = 5;
        for (int ix = 0; ix < numFunctions; ix++) {
            String funcName = "func-" + ix;
            Function.FunctionMetaData func = Function.FunctionMetaData.newBuilder()
                    .setFunctionDetails(Function.FunctionDetails.newBuilder().setName(funcName)
                            .setNamespace("namespace-1").setTenant("tenant-1").setParallelism(1)).setVersion(version)
                    .build();
            functionMetaDataList.add(func);
        }
        doReturn(functionMetaDataList).when(functionMetaDataManager).getAllFunctionMetaData();

        ThreadRuntimeFactory factory = mock(ThreadRuntimeFactory.class);
        doReturn(factory).when(functionRuntimeManager).getRuntimeFactory();

        // set assignments, round-robining the functions among the workers
        int workerIndex = 0;
        List<Function.Assignment> assignmentList = new LinkedList<>();
        for (val func : functionMetaDataList) {
            String workerId = "worker-" + workerIndex;
            Function.Assignment ass = Function.Assignment.newBuilder()
                    .setWorkerId(workerId)
                    .setInstance(Function.Instance.newBuilder()
                            .setFunctionMetaData(func)
                            .setInstanceId(0)
                            .build())
                    .build();
            assignmentList.add(ass);
            workerIndex = (workerIndex + 1) % numWorkers;
        }

        Map<String, Map<String, Function.Assignment>> currentAssignments = new HashMap<>();

        final String workerIdToDrain = "worker-0";
        int numAssignmentsOnDrainedWorker = 0;

        workerIndex = 0;
        for (val ass : assignmentList) {
            String workerId = "worker-" + workerIndex;
            Map<String, Function.Assignment> assignmentEntry = currentAssignments.get(workerId);
            if (assignmentEntry == null) {
                assignmentEntry = new HashMap<>();
            }
            assignmentEntry.put(FunctionCommon.getFullyQualifiedInstanceId(ass.getInstance()), ass);
            currentAssignments.put(workerId, assignmentEntry);
            if (workerId.compareTo(workerIdToDrain) == 0) {
                numAssignmentsOnDrainedWorker++;
            }
            workerIndex = (workerIndex + 1) % numWorkers;
        }
        doReturn(currentAssignments).when(functionRuntimeManager).getCurrentAssignments();

        // i am not leader
        doReturn(false).when(leaderService).isLeader();
        val assignmentsMovedNonLeader = callDrain(workerIdToDrain);
        verify(producer, times(0)).sendAsync(any());
        Assert.assertTrue(assignmentsMovedNonLeader == null);

        // i am leader
        doReturn(true).when(leaderService).isLeader();
        val assignmentsMovedLeader = callDrain(workerIdToDrain);
        List<Invocation> invocations = getMethodInvocationDetails(schedulerManager,
                SchedulerManager.class.getDeclaredMethod("invokeDrain", String.class));
        Assert.assertEquals(invocations.size(), 1);
        verify(errorNotifier, times(0)).triggerError(any());
        Assert.assertFalse(assignmentsMovedLeader.isEmpty());
        // New assignment must contain all the functions.
        Assert.assertEquals(assignmentsMovedLeader.size(), numAssignmentsOnDrainedWorker);
        // Shouldn't see the drained worker in any of the new assignments.
        assignmentsMovedLeader.forEach(ass ->
                Assert.assertTrue(ass.getWorkerId().compareTo(workerIdToDrain) != 0));

        // Get the status of the drain op; since it's a get, the passed-in drain op status is a don't care.
        val dStatus = callGetDrainStatus(workerIdToDrain,
                DrainOps.GetDrainStatus, SchedulerManager.DrainOpStatus.DrainNotInProgress);
        Assert.assertTrue(dStatus.status == LongRunningProcessStatus.Status.SUCCESS);
        Assert.assertTrue(dStatus.lastError.isEmpty());
    }