RabbitMqIO { public static Read read() { return new AutoValue_RabbitMqIO_Read.Builder() .setQueueDeclare(false) .setExchangeDeclare(false) .setMaxReadTime(null) .setMaxNumRecords(Long.MAX_VALUE) .setUseCorrelationId(false) .build(); } private  RabbitMqIO(); static Read read(); static Write write();  }
@Test public void testReadQueue() throws Exception { final int maxNumRecords = 10; PCollection<RabbitMqMessage> raw = p.apply( RabbitMqIO.read() .withUri("amqp: .withQueue("READ") .withMaxNumRecords(maxNumRecords)); PCollection<String> output = raw.apply( MapElements.into(TypeDescriptors.strings()) .via( (RabbitMqMessage message) -> RabbitMqTestUtils.recordToString(message.getBody()))); List<String> records = RabbitMqTestUtils.generateRecords(maxNumRecords).stream() .map(RabbitMqTestUtils::recordToString) .collect(Collectors.toList()); PAssert.that(output).containsInAnyOrder(records); ConnectionFactory connectionFactory = new ConnectionFactory(); connectionFactory.setUri("amqp: Connection connection = null; Channel channel = null; try { connection = connectionFactory.newConnection(); channel = connection.createChannel(); channel.queueDeclare("READ", false, false, false, null); for (String record : records) { channel.basicPublish("", "READ", null, record.getBytes(StandardCharsets.UTF_8)); } p.run(); } finally { if (channel != null) { channel.close(); } if (connection != null) { connection.close(); } } }
@Test public void testReadDeclaredFanoutExchange() throws Exception { doExchangeTest( new ExchangeTestPlan( RabbitMqIO.read().withExchange("DeclaredFanoutExchange", "fanout", "ignored"), 10)); }
@Test public void testReadDeclaredTopicExchangeWithQueueDeclare() throws Exception { doExchangeTest( new ExchangeTestPlan( RabbitMqIO.read() .withExchange("DeclaredTopicExchangeWithQueueDeclare", "topic", "#") .withQueue("declared-queue") .withQueueDeclare(true), 10)); }
@Test public void testReadDeclaredTopicExchange() throws Exception { final int numRecords = 10; RabbitMqIO.Read read = RabbitMqIO.read().withExchange("DeclaredTopicExchange", "topic", "user.create.#"); final Supplier<String> publishRoutingKeyGen = new Supplier<String>() { private AtomicInteger counter = new AtomicInteger(0); @Override public String get() { int count = counter.getAndIncrement(); if (count % 2 == 0) { return "user.create." + count; } return "user.delete." + count; } }; ExchangeTestPlan plan = new ExchangeTestPlan(read, numRecords / 2, numRecords) { @Override public Supplier<String> publishRoutingKeyGen() { return publishRoutingKeyGen; } @Override public List<String> expectedResults() { return IntStream.range(0, numRecords) .filter(i -> i % 2 == 0) .mapToObj(RabbitMqTestUtils::generateRecord) .map(RabbitMqTestUtils::recordToString) .collect(Collectors.toList()); } }; doExchangeTest(plan); }
@Test public void testDeclareIncompatibleExchangeFails() throws Exception { RabbitMqIO.Read read = RabbitMqIO.read().withExchange("IncompatibleExchange", "direct", "unused"); try { doExchangeTest(new ExchangeTestPlan(read, 1), true); fail("Expected to have failed to declare an incompatible exchange"); } catch (Exception e) { Throwable cause = Throwables.getRootCause(e); if (cause instanceof ShutdownSignalException) { ShutdownSignalException sse = (ShutdownSignalException) cause; Method reason = sse.getReason(); if (reason instanceof com.rabbitmq.client.AMQP.Connection.Close) { com.rabbitmq.client.AMQP.Connection.Close close = (com.rabbitmq.client.AMQP.Connection.Close) reason; assertEquals("Expected failure is 530: not-allowed", 530, close.getReplyCode()); } else { fail( "Unexpected ShutdownSignalException reason. Expected Connection.Close. Got: " + reason); } } else { fail("Expected to fail with ShutdownSignalException. Instead failed with " + cause); } } }
@Test public void testUseCorrelationIdSucceedsWhenIdsPresent() throws Exception { int messageCount = 1; AMQP.BasicProperties publishProps = new AMQP.BasicProperties().builder().correlationId("123").build(); doExchangeTest( new ExchangeTestPlan( RabbitMqIO.read() .withExchange("CorrelationIdSuccess", "fanout") .withUseCorrelationId(true), messageCount, messageCount, publishProps)); }
@Test(expected = Pipeline.PipelineExecutionException.class) public void testUseCorrelationIdFailsWhenIdsMissing() throws Exception { int messageCount = 1; AMQP.BasicProperties publishProps = null; doExchangeTest( new ExchangeTestPlan( RabbitMqIO.read() .withExchange("CorrelationIdFailure", "fanout") .withUseCorrelationId(true), messageCount, messageCount, publishProps)); }
@Test(expected = Pipeline.PipelineExecutionException.class) public void testQueueDeclareWithoutQueueNameFails() throws Exception { RabbitMqIO.Read read = RabbitMqIO.read().withQueueDeclare(true); doExchangeTest(new ExchangeTestPlan(read, 1)); }
@Test(timeout = ONE_MINUTE_MS) public void testReadDeclaredFanoutExchange() throws Exception { doExchangeTest( new ExchangeTestPlan( RabbitMqIO.read().withExchange("DeclaredFanoutExchange", "fanout", "ignored"), 10)); }
@Test(timeout = ONE_MINUTE_MS) public void testReadDeclaredTopicExchangeWithQueueDeclare() throws Exception { doExchangeTest( new ExchangeTestPlan( RabbitMqIO.read() .withExchange("DeclaredTopicExchangeWithQueueDeclare", "topic", "#") .withQueue("declared-queue") .withQueueDeclare(true), 10)); }
@Test(timeout = ONE_MINUTE_MS) public void testReadDeclaredTopicExchange() throws Exception { final int numRecords = 10; RabbitMqIO.Read read = RabbitMqIO.read().withExchange("DeclaredTopicExchange", "topic", "user.create.#"); final Supplier<String> publishRoutingKeyGen = new Supplier<String>() { private AtomicInteger counter = new AtomicInteger(0); @Override public String get() { int count = counter.getAndIncrement(); if (count % 2 == 0) { return "user.create." + count; } return "user.delete." + count; } }; ExchangeTestPlan plan = new ExchangeTestPlan(read, numRecords / 2, numRecords) { @Override public Supplier<String> publishRoutingKeyGen() { return publishRoutingKeyGen; } @Override public List<String> expectedResults() { return IntStream.range(0, numRecords) .filter(i -> i % 2 == 0) .mapToObj(RabbitMqTestUtils::generateRecord) .map(RabbitMqTestUtils::recordToString) .collect(Collectors.toList()); } }; doExchangeTest(plan); }
@Test(timeout = ONE_MINUTE_MS) public void testDeclareIncompatibleExchangeFails() throws Exception { RabbitMqIO.Read read = RabbitMqIO.read().withExchange("IncompatibleExchange", "direct", "unused"); try { doExchangeTest(new ExchangeTestPlan(read, 1), true); fail("Expected to have failed to declare an incompatible exchange"); } catch (Exception e) { Throwable cause = Throwables.getRootCause(e); if (cause instanceof ShutdownSignalException) { ShutdownSignalException sse = (ShutdownSignalException) cause; Method reason = sse.getReason(); if (reason instanceof com.rabbitmq.client.AMQP.Connection.Close) { com.rabbitmq.client.AMQP.Connection.Close close = (com.rabbitmq.client.AMQP.Connection.Close) reason; assertEquals("Expected failure is 530: not-allowed", 530, close.getReplyCode()); } else { fail( "Unexpected ShutdownSignalException reason. Expected Connection.Close. Got: " + reason); } } else { fail("Expected to fail with ShutdownSignalException. Instead failed with " + cause); } } }
@Test(timeout = ONE_MINUTE_MS) public void testUseCorrelationIdSucceedsWhenIdsPresent() throws Exception { int messageCount = 1; AMQP.BasicProperties publishProps = new AMQP.BasicProperties().builder().correlationId("123").build(); doExchangeTest( new ExchangeTestPlan( RabbitMqIO.read() .withExchange("CorrelationIdSuccess", "fanout") .withUseCorrelationId(true), messageCount, messageCount, publishProps)); }