Examples with LivyClient com.thinkbiganalytics.kylo.spark.client.LivyClient used on opensource projects

Example 1 with LivyClient

use of com.thinkbiganalytics.kylo.spark.client.LivyClient in project kylo by Teradata.

the class TestLivyHeartbeatMonitor method livyHeartbeatMonitor.

public LivyHeartbeatMonitor livyHeartbeatMonitor() {
    LivyClient mockLivyClient = Mockito.mock(LivyClient.class);
    Session sessionNotStarted = new Session.Builder().id(1).state(SessionState.not_started).build();
    Session sessionStarting = new Session.Builder().id(1).state(SessionState.starting).build();
    Session sessionIdle = new Session.Builder().id(1).state(SessionState.idle).build();
    Session sessioShuttingDown = new Session.Builder().id(1).state(SessionState.shutting_down).build();
    final List<Session> answers = Lists.newArrayList(sessionNotStarted, sessionStarting, sessionIdle, sessioShuttingDown);
    final AtomicInteger numResponse = new AtomicInteger(0);
    JerseyRestClient client = Mockito.mock(JerseyRestClient.class);
    Mockito.when(client.get(Mockito.anyString(), Mockito.eq(Session.class))).thenAnswer(new Answer<Session>() {

        @Override
        public Session answer(InvocationOnMock invocation) throws Throwable {
            try {
                if (numResponse.get() == 3 || numResponse.get() == 8) {
                    // third response and second to last response is session not_found
                    throw new WebApplicationException("Can't find session", 404);
                }
                // Now get from our list of Sessions with certain states
                Session session = answers.remove(0);
                logger.debug("Returning mock response for session with id='{}' and state='{}'", session.getId(), session.getState());
                return session;
            } catch (IndexOutOfBoundsException e) {
                // used up our list of known responses, pretend the server doesn't know the session anymore
                throw new WebApplicationException("Can't find session", 404);
            }
        }
    });
    Mockito.when(mockLivyClient.getSession(Mockito.any(), Mockito.any(SparkLivyProcess.class))).thenAnswer(invocation -> {
        int responseNum = numResponse.addAndGet(1);
        logger.debug("Number of responses return from mockLivyClient:: numResponse={}", responseNum);
        if (responseNum == 1 || responseNum >= 9) {
            // first response and last response is server not started
            throw new ProcessingException(new SocketTimeoutException("Server not started"));
        }
        if (responseNum == 2) {
            // second response is server not HTTPS
            throw new ProcessingException(new SSLHandshakeException("Server is not HTTPS"));
        }
        return livyClient.getSession(client, (SparkLivyProcess) invocation.getArguments()[1]);
    });
    return new LivyHeartbeatMonitor(mockLivyClient, client, livyServer, livyProperties);
}