Examples with TransformRequest com.thinkbiganalytics.spark.rest.model.TransformRequest used on opensource projects

Example 16 with TransformRequest

use of com.thinkbiganalytics.spark.rest.model.TransformRequest in project kylo by Teradata.

the class FileMetadataTest method setup.

private void setup() {
    this.trackerService = Mockito.spy(new FileMetadataTaskService());
    this.sparkShellUserProcessService = new MockSparkShellUserProcessService("tester");
    trackerService.setSparkShellUserProcessService(this.sparkShellUserProcessService);
    restClient = Mockito.mock(SparkShellRestClient.class);
    trackerService.setRestClient(restClient);
    Mockito.doReturn("tester").when(this.trackerService).getUsername();
    // Mockito.doCallRealMethod().when(this.trackerService).runScript(Mockito.any(ModifiedTransformResponse.class),Mockito.any());
    // Mockito.doCallRealMethod().when(this.trackerService).get(Mockito.anyString());
    // Mockito.doCallRealMethod().when(this.trackerService).removeFromCache(Mockito.anyString());
    Mockito.when(restClient.transform(Mockito.any(), Mockito.any())).then(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            TransformRequest r = invocationOnMock.getArgumentAt(1, TransformRequest.class);
            TransformResponse response = new TransformResponse();
            response.setStatus(TransformResponse.Status.PENDING);
            response.setTable(UUID.randomUUID().toString());
            requestTracker.put(response.getTable(), new AtomicInteger(3));
            return response;
        }
    });
    Mockito.when(restClient.getTransformResult(Mockito.any(), Mockito.any())).then(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            String tableId = invocationOnMock.getArgumentAt(1, String.class);
            AtomicInteger i = requestTracker.get(tableId);
            if (i != null) {
                int i1 = i.decrementAndGet();
                if (i1 == 0) {
                    TransformResponse response = new TransformResponse();
                    response.setStatus(TransformResponse.Status.SUCCESS);
                    response.setTable(tableId);
                    response.setResults(headerResult());
                    return Optional.of(response);
                } else {
                    // return pending
                    TransformResponse response = new TransformResponse();
                    response.setTable(tableId);
                    response.setStatus(TransformResponse.Status.PENDING);
                    return Optional.of(response);
                }
            }
            FileMetadataCompletionTask finalResult = trackerService.get(tableId);
            if (finalResult != null) {
                return Optional.of(finalResult.getModifiedTransformResponse());
            } else {
                // return pending
                TransformResponse response = new TransformResponse();
                response.setStatus(TransformResponse.Status.PENDING);
                response.setTable(tableId);
                requestTracker.put(response.getTable(), new AtomicInteger(3));
                return Optional.of(response);
            }
        }
    });
}