Skip to content

Deflake JmsIO tests #39571

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Abacn wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Deflake JmsIO tests #39571

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 6
"modification": 7
}
1 change: 1 addition & 0 deletions CHANGES.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@
* Upgraded Iceberg dependency to 1.11.0 (Java) ([#38925](https://github.com/apache/beam/issues/38925)).
* Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
* Add ArrowFlight IO (Java) ([#20116](https://github.com/apache/beam/issues/20116)).
* (Python) JmsIO (IBM MQ, ActiveMQ, and other providers) is now supported in Python via cross-language ([#30716](https://github.com/apache/beam/issues/30716)).

## New Features / Improvements

Expand Down
176 changes: 3 additions & 173 deletions sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,27 +32,20 @@
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.hamcrest.Matchers.hasProperty;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.isA;
import static org.hamcrest.Matchers.lessThan;
import static org.hamcrest.core.StringContains.containsString;
import static org.hamcrest.object.HasToString.hasToString;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;

import java.io.IOException;
import java.io.NotSerializableException;
import java.io.Serializable;
import java.lang.reflect.Proxy;
import java.nio.ByteBuffer;
Expand All @@ -66,8 +59,6 @@
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import javax.jms.BytesMessage;
Expand All @@ -84,7 +75,6 @@
import org.apache.activemq.command.ActiveMQMessage;
import org.apache.activemq.util.Callback;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.UnboundedSource;
Expand All @@ -94,16 +84,12 @@
import org.apache.beam.sdk.metrics.MetricNameFilter;
import org.apache.beam.sdk.metrics.MetricQueryResults;
import org.apache.beam.sdk.metrics.MetricsFilter;
import org.apache.beam.sdk.options.ExecutorOptions;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.testing.CoderProperties;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.SerializableBiFunction;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
import org.apache.qpid.jms.JmsAcknowledgeCallback;
Expand All @@ -117,7 +103,6 @@
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand Down Expand Up @@ -163,9 +148,9 @@ public static Collection<Object[]> connectionFactories() {
private ConnectionFactory connectionFactory;
private final Class<? extends ConnectionFactory> connectionFactoryClass;
private ConnectionFactory connectionFactoryWithSyncAcksAndWithoutPrefetch;
private final String brokerUrl;
private final Integer brokerPort;
private final String forceAsyncAcksParam;
private String brokerUrl;
private String forceAsyncAcksParam;
private int brokerPort;

public JmsIOTest(
String brokerUrl,
Expand Down Expand Up @@ -252,20 +237,6 @@ public void testReadMessages() throws Exception {
assertQueueIsEmpty();
}

@Test
public void testPipelineWithNonSerializableCF() {
SerializableUtils.ensureSerializable(
JmsIO.read()
.withConnectionFactoryProviderFn(__ -> new MockNonSerializableConnectionFactory()));
try {
SerializableUtils.ensureSerializable(
JmsIO.read().withConnectionFactory(new MockNonSerializableConnectionFactory()));
fail();
} catch (Exception e) {
assertThat(Throwables.getRootCause(e), isA(NotSerializableException.class));
}
}

@Test
public void testReadMessagesWithCFProviderFn() throws Exception {
long count = 5;
Expand Down Expand Up @@ -522,32 +493,6 @@ public void testWriteDynamicMessage() throws Exception {
assertEquals(100, count);
}

@Test
public void testSplitForQueue() throws Exception {
JmsIO.Read read = JmsIO.read().withQueue(QUEUE);
PipelineOptions pipelineOptions = PipelineOptionsFactory.create();
int desiredNumSplits = 5;
JmsIO.UnboundedJmsSource initialSource = new JmsIO.UnboundedJmsSource(read);
List<JmsIO.UnboundedJmsSource> splits = initialSource.split(desiredNumSplits, pipelineOptions);
// in the case of a queue, we have concurrent consumers by default, so the initial number
// splits is equal to the desired number of splits
assertEquals(desiredNumSplits, splits.size());
}

@Test
public void testSplitForTopic() throws Exception {
JmsIO.Read read = JmsIO.read().withTopic(TOPIC);
PipelineOptions pipelineOptions = PipelineOptionsFactory.create();
int desiredNumSplits = 5;
JmsIO.UnboundedJmsSource initialSource = new JmsIO.UnboundedJmsSource(read);
List<JmsIO.UnboundedJmsSource> splits = initialSource.split(desiredNumSplits, pipelineOptions);
// in the case of a topic, we can have only a unique subscriber on the topic per pipeline
// else it means we can have duplicate messages (all subscribers on the topic receive every
// message).
// So, whatever the desizedNumSplits is, the actual number of splits should be 1.
assertEquals(1, splits.size());
}

private boolean advanceWithRetry(UnboundedSource.UnboundedReader reader) throws IOException {
for (int attempt = 0; attempt < 10; attempt++) {
if (reader.advance()) {
Expand Down Expand Up @@ -685,63 +630,6 @@ public void testCheckpointMarkAndFinalizeSeparatelyClientAcknowledgeUnsafe() thr
assertEquals(5, count(QUEUE));
}

@Test
public void testJmsCheckpointMarkIndividualAcknowledgeAllMessages() throws Exception {
Message msg1 = Mockito.mock(Message.class);
Message msg2 = Mockito.mock(Message.class);
Message msg3 = Mockito.mock(Message.class);

JmsCheckpointMark.Preparer preparer =
JmsCheckpointMark.newPreparer(JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE);
preparer.add(msg1);
preparer.add(msg2);
preparer.add(msg3);

AtomicInteger activeCheckpoints = new AtomicInteger(0);
JmsCheckpointMark mark =
preparer.newCheckpoint(
null, null, JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE, activeCheckpoints);
assertNotNull(mark.getMessages());
assertEquals(3, mark.getMessages().size());
assertNull(mark.getConsumer());
assertNull(mark.getSession());
assertEquals(1, activeCheckpoints.get());

mark.finalizeCheckpoint();

Mockito.verify(msg1, Mockito.times(1)).acknowledge();
Mockito.verify(msg2, Mockito.times(1)).acknowledge();
Mockito.verify(msg3, Mockito.times(1)).acknowledge();
assertEquals(0, activeCheckpoints.get());
}

@Test
public void testJmsCheckpointMarkClientAcknowledgeUnsafeNoSessionRecreation() throws Exception {
Message msg1 = Mockito.mock(Message.class);
Message msg2 = Mockito.mock(Message.class);

JmsCheckpointMark.Preparer preparer =
JmsCheckpointMark.newPreparer(JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE_UNSAFE);
preparer.add(msg1);
preparer.add(msg2);

AtomicInteger activeCheckpoints = new AtomicInteger(0);
JmsCheckpointMark mark =
preparer.newCheckpoint(
null, null, JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE_UNSAFE, activeCheckpoints);
assertNotNull(mark.getMessages());
assertEquals(1, mark.getMessages().size());
assertNull(mark.getConsumer());
assertNull(mark.getSession());
assertEquals(1, activeCheckpoints.get());

mark.finalizeCheckpoint();

Mockito.verify(msg2, Mockito.times(1)).acknowledge();
Mockito.verify(msg1, Mockito.never()).acknowledge();
assertEquals(0, activeCheckpoints.get());
}

private JmsIO.UnboundedJmsReader setupReaderForTest() throws JMSException {
return setupReaderForTest(null);
}
Expand Down Expand Up @@ -890,17 +778,6 @@ public void testCheckpointMarkSafety() throws Exception {
runner.join();
}

/** Test the checkpoint mark default coder, which is actually AvroCoder. */
@Test
public void testCheckpointMarkDefaultCoder() throws Exception {
JmsCheckpointMark jmsCheckpointMark =
JmsCheckpointMark.newPreparer(JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE)
.newCheckpoint(null, null, JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE, null);
Coder coder = new JmsIO.UnboundedJmsSource(null).getCheckpointMarkCoder();
CoderProperties.coderSerializable(coder);
CoderProperties.coderDecodeEncodeEqual(coder, jmsCheckpointMark);
}

@Test
public void testDefaultAutoscaler() throws IOException {
JmsIO.Read spec =
Expand Down Expand Up @@ -945,53 +822,6 @@ public void testCustomAutoscaler() throws IOException {
verify(autoScaler, times(1)).stop();
}

@Test
public void testCloseWithTimeout() throws IOException, JMSException {
Duration closeTimeout = Duration.millis(2000L);
JmsIO.Read spec =
JmsIO.read()
.withConnectionFactory(connectionFactory)
.withUsername(USERNAME)
.withPassword(PASSWORD)
.withQueue(QUEUE)
.withCloseTimeout(closeTimeout);

JmsIO.UnboundedJmsSource source = new JmsIO.UnboundedJmsSource(spec);

ScheduledExecutorService mockScheduledExecutorService =
Mockito.mock(ScheduledExecutorService.class);
ExecutorOptions options = PipelineOptionsFactory.as(ExecutorOptions.class);
options.setScheduledExecutorService(mockScheduledExecutorService);
ArgumentCaptor<Runnable> runnableArgumentCaptor = ArgumentCaptor.forClass(Runnable.class);
when(mockScheduledExecutorService.schedule(
runnableArgumentCaptor.capture(), anyLong(), any(TimeUnit.class)))
.thenReturn(null /* unused */);

JmsIO.UnboundedJmsReader reader = source.createReader(options, null);
reader.start();
assertFalse(getDiscardedValue(reader));
reader.checkpointMarkPreparer.add(Mockito.mock(Message.class));
CheckpointMark mark = reader.getCheckpointMark();
reader.close();
assertTrue(getDiscardedValue(reader));
verify(mockScheduledExecutorService)
.schedule(any(Runnable.class), eq(1L), eq(TimeUnit.SECONDS));
mark.finalizeCheckpoint();
runnableArgumentCaptor.getValue().run();
assertTrue(getDiscardedValue(reader));
verifyNoMoreInteractions(mockScheduledExecutorService);
}

private boolean getDiscardedValue(JmsIO.UnboundedJmsReader reader) {
JmsCheckpointMark.Preparer preparer = reader.checkpointMarkPreparer;
preparer.lock.readLock().lock();
try {
return preparer.discarded;
} finally {
preparer.lock.readLock().unlock();
}
}

@Test
public void testDiscardCheckpointMark() throws Exception {
Connection connection =
Expand Down
Loading
Loading