SolrIO { public static Write write() { return new AutoValue_SolrIO_Write.Builder().setMaxBatchSize(1000).build(); } private  SolrIO(); static Read read(); static ReadAll readAll(); static Write write();  }
@Test public void testWrite() throws Exception { List<SolrInputDocument> data = SolrIOTestUtils.createDocuments(NUM_DOCS); SolrIO.Write write = SolrIO.write().withConnectionConfiguration(connectionConfiguration).to(SOLR_COLLECTION); pipeline.apply(Create.of(data)).apply(write); pipeline.run(); long currentNumDocs = SolrIOTestUtils.commitAndGetCurrentNumDocs(SOLR_COLLECTION, solrClient); assertEquals(NUM_DOCS, currentNumDocs); QueryResponse response = solrClient.query(SOLR_COLLECTION, new SolrQuery("scientist:Lovelace")); assertEquals(NUM_DOCS / NUM_SCIENTISTS, response.getResults().getNumFound()); }
@Test public void testWriteWithMaxBatchSize() throws Exception { SolrIO.Write write = SolrIO.write() .withConnectionConfiguration(connectionConfiguration) .to(SOLR_COLLECTION) .withMaxBatchSize(BATCH_SIZE); try (DoFnTester<SolrInputDocument, Void> fnTester = DoFnTester.of(new SolrIO.Write.WriteFn(write))) { List<SolrInputDocument> input = SolrIOTestUtils.createDocuments(NUM_DOCS); long numDocsProcessed = 0; long numDocsInserted = 0; for (SolrInputDocument document : input) { fnTester.processElement(document); numDocsProcessed++; if ((numDocsProcessed % 100) == 0) { long currentNumDocs = SolrIOTestUtils.commitAndGetCurrentNumDocs(SOLR_COLLECTION, solrClient); if ((numDocsProcessed % BATCH_SIZE) == 0) { assertEquals( "we are at the end of a bundle, we should have inserted all processed documents", numDocsProcessed, currentNumDocs); numDocsInserted = currentNumDocs; } else { assertEquals( "we are not at the end of a bundle, we should have inserted no more documents", numDocsInserted, currentNumDocs); } } } } }
@Test public void testWriteRetry() throws Throwable { thrown.expect(IOException.class); thrown.expectMessage("Error writing to Solr"); @SuppressWarnings("unchecked") Set<Object> entryState = ImmutableSet.copyOf(ObjectReleaseTracker.OBJECTS.keySet()); SolrIO.Write write = SolrIO.write() .withConnectionConfiguration(connectionConfiguration) .withRetryConfiguration( SolrIO.RetryConfiguration.create(3, Duration.standardMinutes(3)) .withRetryPredicate(new LenientRetryPredicate())) .to("wrong-collection"); List<SolrInputDocument> data = SolrIOTestUtils.createDocuments(NUM_DOCS); pipeline.apply(Create.of(data)).apply(write); try { pipeline.run(); } catch (final Pipeline.PipelineExecutionException e) { int waitAttempts = 30; while (namedThreadIsAlive("direct-runner-worker") && waitAttempts-- >= 0) { LOG.info("Pausing to allow direct-runner-worker threads to finish"); Thread.sleep(1000); } for (Object o : ObjectReleaseTracker.OBJECTS.keySet()) { if (o instanceof SolrZkClient && !entryState.contains(o)) { LOG.info("Removing unreleased SolrZkClient"); ObjectReleaseTracker.release(o); } } expectedLogs.verifyWarn(String.format(SolrIO.Write.WriteFn.RETRY_ATTEMPT_LOG, 1)); expectedLogs.verifyWarn(String.format(SolrIO.Write.WriteFn.RETRY_ATTEMPT_LOG, 2)); throw e.getCause(); } fail("Pipeline should not have run to completion"); }
@Test public void testBatchSize() { SolrIO.Write write1 = SolrIO.write() .withConnectionConfiguration(connectionConfiguration) .withMaxBatchSize(BATCH_SIZE); assertTrue(write1.getMaxBatchSize() == BATCH_SIZE); SolrIO.Write write2 = SolrIO.write().withConnectionConfiguration(connectionConfiguration); assertTrue(write2.getMaxBatchSize() == DEFAULT_BATCH_SIZE); }