use of com.thinkbiganalytics.spark.rest.filemetadata.tasks.FileMetadataTaskService 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);
}
}
});
}
Aggregations