Examples with Extractor org.apache.gobblin.source.extractor.Extractor used on opensource projects

Example 6 with Extractor

use of org.apache.gobblin.source.extractor.Extractor in project incubator-gobblin by apache.

the class GoogleAnalyticsUnsampledExtractorTest method setup.

private GoogleAnalyticsUnsampledExtractor setup(final ReportCreationStatus status, WorkUnitState wuState, final boolean throwException) throws IOException {
    Extractor actualExtractor = mock(Extractor.class);
    gaService = mock(Analytics.class);
    Management mgmt = mock(Management.class);
    when(gaService.management()).thenReturn(mgmt);
    UnsampledReports req = mock(UnsampledReports.class);
    when(mgmt.unsampledReports()).thenReturn(req);
    getReq = mock(Get.class);
    when(req.get(anyString(), anyString(), anyString(), anyString())).thenReturn(getReq);
    int pollCount = 10;
    final MutableInt countDown = new MutableInt(pollCount);
    when(getReq.execute()).then(new Answer<UnsampledReport>() {

        @Override
        public UnsampledReport answer(InvocationOnMock invocation) throws Throwable {
            countDown.decrement();
            if (countDown.intValue() == 0) {
                UnsampledReport response = new UnsampledReport();
                DriveDownloadDetails details = new DriveDownloadDetails();
                details.setDocumentId(EXPECTED_FILE_ID);
                response.setStatus(status.name()).setDownloadType(DOWNLOAD_TYPE_GOOGLE_DRIVE).setDriveDownloadDetails(details);
                return response;
            } else if (throwException) {
                throw new RuntimeException("Dummy exception.");
            }
            return new UnsampledReport();
        }
    });
    return new GoogleAnalyticsUnsampledExtractor<>(wuState, actualExtractor, gaService);
}