Examples with JavaTestKit akka.testkit.JavaTestKit used on opensource projects

Example 41 with JavaTestKit

use of akka.testkit.JavaTestKit in project flink by apache.

the class TaskManagerRegistrationTest method testCheckForValidRegistrationSessionIDs.

@Test
public void testCheckForValidRegistrationSessionIDs() {
    new JavaTestKit(actorSystem) {

        {
            ActorGateway taskManagerGateway = null;
            try {
                // we make the test actor (the test kit) the JobManager to intercept
                // the messages
                taskManagerGateway = createTaskManager(actorSystem, getTestActor(), config, true, false);
                final ActorRef taskManager = taskManagerGateway.actor();
                final UUID falseLeaderSessionID = UUID.randomUUID();
                final UUID trueLeaderSessionID = null;
                new Within(timeout) {

                    @Override
                    protected void run() {
                        taskManager.tell(TaskManagerMessages.getNotifyWhenRegisteredAtJobManagerMessage(), getTestActor());
                        // the TaskManager should try to register
                        LeaderSessionMessage lsm = expectMsgClass(LeaderSessionMessage.class);
                        assertTrue(lsm.leaderSessionID() == trueLeaderSessionID);
                        assertTrue(lsm.message() instanceof RegisterTaskManager);
                        final ActorRef tm = getLastSender();
                        // This AcknowledgeRegistration message should be discarded because the
                        // registration session ID is wrong
                        tm.tell(new LeaderSessionMessage(falseLeaderSessionID, new AcknowledgeRegistration(new InstanceID(), 1)), getTestActor());
                        // Valid AcknowledgeRegistration message
                        tm.tell(new LeaderSessionMessage(trueLeaderSessionID, new AcknowledgeRegistration(new InstanceID(), 1)), getTestActor());
                        Object message = null;
                        Object confirmMessageClass = TaskManagerMessages.getRegisteredAtJobManagerMessage().getClass();
                        while (message == null || !(message.getClass().equals(confirmMessageClass))) {
                            message = receiveOne(TestingUtils.TESTING_DURATION());
                        }
                        tm.tell(JobManagerMessages.getRequestLeaderSessionID(), getTestActor());
                        expectMsgEquals(new JobManagerMessages.ResponseLeaderSessionID(trueLeaderSessionID));
                    }
                };
            } catch (Throwable e) {
                e.printStackTrace();
                fail(e.getMessage());
            } finally {
                stopActor(taskManagerGateway);
            }
        }
    };
}

Example 42 with JavaTestKit

use of akka.testkit.JavaTestKit in project flink by apache.

the class TaskManagerRegistrationTest method testTaskManagerResumesConnectAfterRefusedRegistration.

/**
	 * Make sure that the TaskManager keeps trying to register, even after
	 * registration attempts have been refused.
	 */
@Test
public void testTaskManagerResumesConnectAfterRefusedRegistration() {
    new JavaTestKit(actorSystem) {

        {
            ActorGateway jm = null;
            ActorGateway taskManager = null;
            try {
                jm = TestingUtils.createForwardingActor(actorSystem, getTestActor(), Option.<String>empty());
                final ActorGateway jmGateway = jm;
                FiniteDuration refusedRegistrationPause = new FiniteDuration(500, TimeUnit.MILLISECONDS);
                Configuration tmConfig = new Configuration(config);
                tmConfig.setString(ConfigConstants.TASK_MANAGER_REFUSED_REGISTRATION_PAUSE, refusedRegistrationPause.toString());
                // we make the test actor (the test kit) the JobManager to intercept
                // the messages
                taskManager = createTaskManager(actorSystem, jmGateway, tmConfig, true, false);
                final ActorGateway taskManagerGateway = taskManager;
                // check and decline initial registration
                new Within(timeout) {

                    @Override
                    protected void run() {
                        // the TaskManager should try to register
                        expectMsgClass(RegisterTaskManager.class);
                        // we decline the registration
                        taskManagerGateway.tell(new RefuseRegistration(new Exception("test reason")), jmGateway);
                    }
                };
                // the TaskManager should wait a bit an retry...
                FiniteDuration maxDelay = (FiniteDuration) refusedRegistrationPause.$times(3.0);
                new Within(maxDelay) {

                    @Override
                    protected void run() {
                        expectMsgClass(RegisterTaskManager.class);
                    }
                };
            } catch (Throwable e) {
                e.printStackTrace();
                fail(e.getMessage());
            } finally {
                stopActor(taskManager);
                stopActor(jm);
            }
        }
    };
}

Example 43 with JavaTestKit

use of akka.testkit.JavaTestKit in project flink by apache.

the class TaskManagerTest method testLocalPartitionNotFound.

/**
	 *  Tests that repeated local {@link PartitionNotFoundException}s ultimately fail the receiver.
	 */
@Test
public void testLocalPartitionNotFound() throws Exception {
    new JavaTestKit(system) {

        {
            ActorGateway jobManager = null;
            ActorGateway taskManager = null;
            final ActorGateway testActorGateway = new AkkaActorGateway(getTestActor(), leaderSessionID);
            try {
                final IntermediateDataSetID resultId = new IntermediateDataSetID();
                // Create the JM
                ActorRef jm = system.actorOf(Props.create(new SimplePartitionStateLookupJobManagerCreator(leaderSessionID, getTestActor())));
                jobManager = new AkkaActorGateway(jm, leaderSessionID);
                final Configuration config = new Configuration();
                config.setInteger(TaskManagerOptions.NETWORK_REQUEST_BACKOFF_INITIAL, 100);
                config.setInteger(TaskManagerOptions.NETWORK_REQUEST_BACKOFF_MAX, 200);
                taskManager = TestingUtils.createTaskManager(system, jobManager, config, true, true);
                // ---------------------------------------------------------------------------------
                final ActorGateway tm = taskManager;
                final JobID jid = new JobID();
                final JobVertexID vid = new JobVertexID();
                final ExecutionAttemptID eid = new ExecutionAttemptID();
                final ResultPartitionID partitionId = new ResultPartitionID();
                // Local location (on the same TM though) for the partition
                final ResultPartitionLocation loc = ResultPartitionLocation.createLocal();
                final InputChannelDeploymentDescriptor[] icdd = new InputChannelDeploymentDescriptor[] { new InputChannelDeploymentDescriptor(partitionId, loc) };
                final InputGateDeploymentDescriptor igdd = new InputGateDeploymentDescriptor(resultId, ResultPartitionType.PIPELINED, 0, icdd);
                final TaskDeploymentDescriptor tdd = createTaskDeploymentDescriptor(jid, "TestJob", vid, eid, new SerializedValue<>(new ExecutionConfig()), "Receiver", 1, 0, 1, 0, new Configuration(), new Configuration(), Tasks.AgnosticReceiver.class.getName(), Collections.<ResultPartitionDeploymentDescriptor>emptyList(), Collections.singletonList(igdd), Collections.<BlobKey>emptyList(), Collections.<URL>emptyList(), 0);
                new Within(new FiniteDuration(120, TimeUnit.SECONDS)) {

                    @Override
                    protected void run() {
                        // Submit the task
                        tm.tell(new SubmitTask(tdd), testActorGateway);
                        expectMsgClass(Acknowledge.get().getClass());
                        // Wait to be notified about the final execution state by the mock JM
                        TaskExecutionState msg = expectMsgClass(TaskExecutionState.class);
                        // The task should fail after repeated requests
                        assertEquals(msg.getExecutionState(), ExecutionState.FAILED);
                        Throwable error = msg.getError(getClass().getClassLoader());
                        if (error.getClass() != PartitionNotFoundException.class) {
                            error.printStackTrace();
                            fail("Wrong exception: " + error.getMessage());
                        }
                    }
                };
            } catch (Exception e) {
                e.printStackTrace();
                fail(e.getMessage());
            } finally {
                TestingUtils.stopActor(taskManager);
                TestingUtils.stopActor(jobManager);
            }
        }
    };
}

Example 44 with JavaTestKit

use of akka.testkit.JavaTestKit in project flink by apache.

the class AccumulatorLiveITCase method verifyResults.

private static void verifyResults() {
    new JavaTestKit(system) {

        {
            ActorGateway selfGateway = new AkkaActorGateway(getRef(), jobManagerGateway.leaderSessionID());
            // register for accumulator changes
            jobManagerGateway.tell(new TestingJobManagerMessages.NotifyWhenAccumulatorChange(jobID), selfGateway);
            expectMsgEquals(TIMEOUT, true);
            // submit job
            jobManagerGateway.tell(new JobManagerMessages.SubmitJob(jobGraph, ListeningBehaviour.EXECUTION_RESULT), selfGateway);
            expectMsgClass(TIMEOUT, JobManagerMessages.JobSubmitSuccess.class);
            TestingJobManagerMessages.UpdatedAccumulators msg = (TestingJobManagerMessages.UpdatedAccumulators) receiveOne(TIMEOUT);
            Map<String, Accumulator<?, ?>> userAccumulators = msg.userAccumulators();
            ExecutionAttemptID mapperTaskID = null;
            ExecutionAttemptID sinkTaskID = null;
            /* Check for accumulator values */
            if (checkUserAccumulators(0, userAccumulators)) {
                LOG.info("Passed initial check for map task.");
            } else {
                fail("Wrong accumulator results when map task begins execution.");
            }
            int expectedAccVal = 0;
            /* for mapper task */
            for (int i = 1; i <= NUM_ITERATIONS; i++) {
                expectedAccVal += i;
                // receive message
                msg = (TestingJobManagerMessages.UpdatedAccumulators) receiveOne(TIMEOUT);
                userAccumulators = msg.userAccumulators();
                LOG.info("{}", userAccumulators);
                if (checkUserAccumulators(expectedAccVal, userAccumulators)) {
                    LOG.info("Passed round #" + i);
                } else if (checkUserAccumulators(expectedAccVal, userAccumulators)) {
                    // we determined the wrong task id and need to switch the two here
                    ExecutionAttemptID temp = mapperTaskID;
                    mapperTaskID = sinkTaskID;
                    sinkTaskID = temp;
                    LOG.info("Passed round #" + i);
                } else {
                    fail("Failed in round #" + i);
                }
            }
            msg = (TestingJobManagerMessages.UpdatedAccumulators) receiveOne(TIMEOUT);
            userAccumulators = msg.userAccumulators();
            if (checkUserAccumulators(expectedAccVal, userAccumulators)) {
                LOG.info("Passed initial check for sink task.");
            } else {
                fail("Wrong accumulator results when sink task begins execution.");
            }
            /* for sink task */
            for (int i = 1; i <= NUM_ITERATIONS; i++) {
                // receive message
                msg = (TestingJobManagerMessages.UpdatedAccumulators) receiveOne(TIMEOUT);
                userAccumulators = msg.userAccumulators();
                LOG.info("{}", userAccumulators);
                if (checkUserAccumulators(expectedAccVal, userAccumulators)) {
                    LOG.info("Passed round #" + i);
                } else {
                    fail("Failed in round #" + i);
                }
            }
            expectMsgClass(TIMEOUT, JobManagerMessages.JobResultSuccess.class);
        }
    };
}

Example 45 with JavaTestKit

use of akka.testkit.JavaTestKit in project flink by apache.

the class UtilsTest method testYarnFlinkResourceManagerJobManagerLostLeadership.

@Test
public void testYarnFlinkResourceManagerJobManagerLostLeadership() throws Exception {
    new JavaTestKit(system) {

        {
            final Deadline deadline = new FiniteDuration(3, TimeUnit.MINUTES).fromNow();
            Configuration flinkConfig = new Configuration();
            YarnConfiguration yarnConfig = new YarnConfiguration();
            TestingLeaderRetrievalService leaderRetrievalService = new TestingLeaderRetrievalService();
            String applicationMasterHostName = "localhost";
            String webInterfaceURL = "foobar";
            ContaineredTaskManagerParameters taskManagerParameters = new ContaineredTaskManagerParameters(1l, 1l, 1l, 1, new HashMap<String, String>());
            ContainerLaunchContext taskManagerLaunchContext = mock(ContainerLaunchContext.class);
            int yarnHeartbeatIntervalMillis = 1000;
            int maxFailedContainers = 10;
            int numInitialTaskManagers = 5;
            final YarnResourceManagerCallbackHandler callbackHandler = new YarnResourceManagerCallbackHandler();
            AMRMClientAsync<AMRMClient.ContainerRequest> resourceManagerClient = mock(AMRMClientAsync.class);
            NMClient nodeManagerClient = mock(NMClient.class);
            UUID leaderSessionID = UUID.randomUUID();
            final List<Container> containerList = new ArrayList<>();
            for (int i = 0; i < numInitialTaskManagers; i++) {
                containerList.add(new TestingContainer("container_" + i, "localhost"));
            }
            doAnswer(new Answer() {

                int counter = 0;

                @Override
                public Object answer(InvocationOnMock invocation) throws Throwable {
                    if (counter < containerList.size()) {
                        callbackHandler.onContainersAllocated(Collections.singletonList(containerList.get(counter++)));
                    }
                    return null;
                }
            }).when(resourceManagerClient).addContainerRequest(Matchers.any(AMRMClient.ContainerRequest.class));
            ActorRef resourceManager = null;
            ActorRef leader1;
            try {
                leader1 = system.actorOf(Props.create(TestingUtils.ForwardingActor.class, getRef(), Option.apply(leaderSessionID)));
                resourceManager = system.actorOf(Props.create(TestingYarnFlinkResourceManager.class, flinkConfig, yarnConfig, leaderRetrievalService, applicationMasterHostName, webInterfaceURL, taskManagerParameters, taskManagerLaunchContext, yarnHeartbeatIntervalMillis, maxFailedContainers, numInitialTaskManagers, callbackHandler, resourceManagerClient, nodeManagerClient));
                leaderRetrievalService.notifyListener(leader1.path().toString(), leaderSessionID);
                final AkkaActorGateway leader1Gateway = new AkkaActorGateway(leader1, leaderSessionID);
                final AkkaActorGateway resourceManagerGateway = new AkkaActorGateway(resourceManager, leaderSessionID);
                doAnswer(new Answer() {

                    @Override
                    public Object answer(InvocationOnMock invocation) throws Throwable {
                        Container container = (Container) invocation.getArguments()[0];
                        resourceManagerGateway.tell(new NotifyResourceStarted(YarnFlinkResourceManager.extractResourceID(container)), leader1Gateway);
                        return null;
                    }
                }).when(nodeManagerClient).startContainer(Matchers.any(Container.class), Matchers.any(ContainerLaunchContext.class));
                expectMsgClass(deadline.timeLeft(), RegisterResourceManager.class);
                resourceManagerGateway.tell(new RegisterResourceManagerSuccessful(leader1, Collections.EMPTY_LIST));
                for (int i = 0; i < containerList.size(); i++) {
                    expectMsgClass(deadline.timeLeft(), Acknowledge.class);
                }
                Future<Object> taskManagerRegisteredFuture = resourceManagerGateway.ask(new NotifyWhenResourcesRegistered(numInitialTaskManagers), deadline.timeLeft());
                Await.ready(taskManagerRegisteredFuture, deadline.timeLeft());
                leaderRetrievalService.notifyListener(null, null);
                leaderRetrievalService.notifyListener(leader1.path().toString(), leaderSessionID);
                expectMsgClass(deadline.timeLeft(), RegisterResourceManager.class);
                resourceManagerGateway.tell(new RegisterResourceManagerSuccessful(leader1, Collections.EMPTY_LIST));
                for (Container container : containerList) {
                    resourceManagerGateway.tell(new NotifyResourceStarted(YarnFlinkResourceManager.extractResourceID(container)), leader1Gateway);
                }
                for (int i = 0; i < containerList.size(); i++) {
                    expectMsgClass(deadline.timeLeft(), Acknowledge.class);
                }
                Future<Object> numberOfRegisteredResourcesFuture = resourceManagerGateway.ask(RequestNumberOfRegisteredResources.Instance, deadline.timeLeft());
                int numberOfRegisteredResources = (Integer) Await.result(numberOfRegisteredResourcesFuture, deadline.timeLeft());
                assertEquals(numInitialTaskManagers, numberOfRegisteredResources);
            } finally {
                if (resourceManager != null) {
                    resourceManager.tell(PoisonPill.getInstance(), ActorRef.noSender());
                }
            }
        }
    };
}