Examples with BulkBatchResult org.teiid.translator.salesforce.SalesforceConnection.BulkBatchResult used on opensource projects

Example 1 with BulkBatchResult

use of org.teiid.translator.salesforce.SalesforceConnection.BulkBatchResult in project teiid by teiid.

the class TestQueryExecutionImpl method testBulkFlow.

@Test
public void testBulkFlow() throws Exception {
    // $NON-NLS-1$
    Select command = (Select) translationUtility.parseCommand("select Name from Account");
    SalesforceConnection connection = Mockito.mock(SalesforceConnection.class);
    JobInfo jobInfo = Mockito.mock(JobInfo.class);
    Mockito.when(connection.createBulkJob(Mockito.anyString(), Mockito.eq(OperationEnum.query), Mockito.eq(true))).thenReturn(jobInfo);
    final BatchResultInfo info = new BatchResultInfo("x");
    Mockito.when(connection.getBatchQueryResults(Mockito.anyString(), Mockito.eq(info))).thenAnswer(new Answer<BulkBatchResult>() {

        boolean first = true;

        @Override
        public BulkBatchResult answer(InvocationOnMock invocation) throws Throwable {
            if (first) {
                first = false;
                throw new DataNotAvailableException();
            }
            if (info.getAndIncrementResultNum() == 0) {
                final Iterator<List<String>> i = Arrays.asList(Arrays.asList("Name"), Arrays.asList("X")).iterator();
                return new BulkBatchResult() {

                    @Override
                    public List<String> nextRecord() throws IOException {
                        if (!i.hasNext()) {
                            return null;
                        }
                        return i.next();
                    }

                    @Override
                    public void close() {
                    }
                };
            }
            return null;
        }
    });
    Mockito.when(connection.addBatch("SELECT Account.Name FROM Account", jobInfo)).thenReturn(info);
    ExecutionContext mock = Mockito.mock(ExecutionContext.class);
    Mockito.stub(mock.getSourceHints()).toReturn(Arrays.asList("bulk"));
    QueryExecutionImpl execution = new QueryExecutionImpl(command, connection, Mockito.mock(RuntimeMetadata.class), mock, new SalesForceExecutionFactory());
    execution.execute();
    try {
        execution.next();
        fail();
    } catch (DataNotAvailableException e) {
    }
    List<?> row = execution.next();
    assertEquals(Arrays.asList("X"), row);
    assertNull(execution.next());
}