Fix Flaky JmsIOTest

Author: shnapz

sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java

@@ -648,6 +648,9 @@ private JmsIO.UnboundedJmsReader setupReaderForTest() throws JMSException {
     final int delay = 10;
     return connectorClass == JmsConnectionFactory.class
         ? (JmsTextMessage message) -> {
+          if (message == null) {
+            return null;
+          }
           final JmsAcknowledgeCallback originalCallback = message.getAcknowledgeCallback();
           JmsAcknowledgeCallback jmsAcknowledgeCallbackMock =
               Mockito.mock(JmsAcknowledgeCallback.class);
@@ -667,6 +670,9 @@ private JmsIO.UnboundedJmsReader setupReaderForTest() throws JMSException {
           return message;
         }
         : (ActiveMQMessage message) -> {
+          if (message == null) {
+            return null;
+          }
           final Callback originalCallback = message.getAcknowledgeCallback();
           message.setAcknowledgeCallback(
               () -> {