Migrated and fixed DiscoveryNetworkReconnectTest (it was not working before the migration)

Author: kenliao94

activemq-unit-tests/src/test/java/org/apache/activemq/transport/discovery/DiscoveryNetworkReconnectTest.java

@@ -17,6 +17,13 @@
 package org.apache.activemq.transport.discovery;
 
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import java.net.URI;
 import java.util.concurrent.Semaphore;
@@ -30,24 +37,13 @@
 import org.apache.activemq.transport.discovery.multicast.MulticastDiscoveryAgentFactory;
 import org.apache.activemq.util.SocketProxy;
 import org.apache.activemq.util.Wait;
-import org.hamcrest.BaseMatcher;
-import org.hamcrest.Description;
-import org.jmock.Expectations;
-import org.jmock.Mockery;
-import org.jmock.api.Invocation;
-import org.jmock.integration.junit4.JMock;
-import org.jmock.integration.junit4.JUnit4Mockery;
-import org.jmock.lib.action.CustomAction;
-import org.jmock.lib.legacy.ClassImposteriser;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
-import org.junit.runner.RunWith;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 
-@RunWith(JMock.class)
 public class DiscoveryNetworkReconnectTest {
 
     private static final Logger LOG = LoggerFactory.getLogger(DiscoveryNetworkReconnectTest.class);
@@ -57,42 +53,12 @@ public class DiscoveryNetworkReconnectTest {
     final Semaphore mbeanRegistered = new Semaphore(0);
     final Semaphore mbeanUnregistered = new Semaphore(0);
     BrokerService brokerA, brokerB;
-    Mockery context;
     ManagementContext managementContext;
     DiscoveryAgent agent;
     SocketProxy proxy;
 
-    // ignore the hostname resolution component as this is machine dependent
-    class NetworkBridgeObjectNameMatcher<T> extends BaseMatcher<T> {
-        T name;
-        NetworkBridgeObjectNameMatcher(T o) {
-            name = o;
-        }
-
-        @Override
-        public boolean matches(Object arg0) {
-            ObjectName other = (ObjectName) arg0;
-            ObjectName mine = (ObjectName) name;
-            LOG.info("Match: " + mine + " vs: " + other);
-
-            if (!"networkConnectors".equals(other.getKeyProperty("connector"))) {
-                return false;
-            }
-            return other.getKeyProperty("connector").equals(mine.getKeyProperty("connector")) &&
-                   other.getKeyProperty("networkBridge") != null && mine.getKeyProperty("networkBridge") != null;
-        }
-
-        @Override
-        public void describeTo(Description arg0) {
-            arg0.appendText(this.getClass().getName());
-        }
-    }
-
     @Before
     public void setUp() throws Exception {
-        context = new JUnit4Mockery() {{
-            setImposteriser(ClassImposteriser.INSTANCE);
-        }};
         brokerA = new BrokerService();
         brokerA.setBrokerName("BrokerA");
         configure(brokerA);
@@ -101,73 +67,31 @@ public void setUp() throws Exception {
         brokerA.waitUntilStarted();
 
         proxy = new SocketProxy(brokerA.getTransportConnectors().get(0).getConnectUri());
-        managementContext = context.mock(ManagementContext.class);
-
-        context.checking(new Expectations(){{
-            allowing(managementContext).getJmxDomainName(); will (returnValue("Test"));
-            allowing(managementContext).setBrokerName("BrokerNC");
-            allowing(managementContext).start();
-            allowing(managementContext).isCreateConnector();
-            allowing(managementContext).stop();
-            allowing(managementContext).isConnectorStarted();
-
-            // expected MBeans
-            allowing(managementContext).registerMBean(with(any(Object.class)), with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC"))));
-            allowing(managementContext).registerMBean(with(any(Object.class)), with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,service=Health"))));
-            allowing(managementContext).registerMBean(with(any(Object.class)), with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,connector=networkConnectors,networkConnectorName=NC"))));
-            allowing(managementContext).registerMBean(with(any(Object.class)), with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,service=Log4JConfiguration"))));
-            allowing(managementContext).registerMBean(with(any(Object.class)), with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,destinationType=Topic,destinationName=ActiveMQ.Advisory.Connection"))));
-            allowing(managementContext).registerMBean(with(any(Object.class)), with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,destinationType=Topic,destinationName=ActiveMQ.Advisory.NetworkBridge"))));
-            allowing(managementContext).registerMBean(with(any(Object.class)), with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,destinationType=Topic,destinationName=ActiveMQ.Advisory.MasterBroker"))));
-            allowing(managementContext).registerMBean(with(any(Object.class)), with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,service=jobScheduler,jobSchedulerName=JMS"))));
-            allowing(managementContext).getObjectInstance(with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,connector=networkConnectors,networkConnectorName=NC"))));
-
-
-            atLeast(maxReconnects - 1).of (managementContext).registerMBean(with(any(Object.class)), with(new NetworkBridgeObjectNameMatcher<ObjectName>(
-                        new ObjectName("Test:type=Broker,brokerName=BrokerNC,connector=networkConnectors,networkConnectorName=NC,networkBridge=localhost/127.0.0.1_"
-                            + proxy.getUrl().getPort())))); will(new CustomAction("signal register network mbean") {
-                                @Override
-                                public Object invoke(Invocation invocation) throws Throwable {
-                                    LOG.info("Mbean Registered: " + invocation.getParameter(0));
-                                    mbeanRegistered.release();
-                                    return new ObjectInstance((ObjectName)invocation.getParameter(1), "discription");
-                                }
-                            });
-            atLeast(maxReconnects - 1).of (managementContext).unregisterMBean(with(new NetworkBridgeObjectNameMatcher<ObjectName>(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,connector=networkConnectors,networkConnectorName=NC,networkBridge=localhost/127.0.0.1_"
-                            + proxy.getUrl().getPort())))); will(new CustomAction("signal unregister network mbean") {
-                                @Override
-                                public Object invoke(Invocation invocation) throws Throwable {
-                                    LOG.info("Mbean Unregistered: " + invocation.getParameter(0));
-                                    mbeanUnregistered.release();
-                                    return null;
-                                }
-                            });
-
-            allowing(managementContext).unregisterMBean(with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC"))));
-            allowing(managementContext).unregisterMBean(with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,service=Health"))));
-            allowing(managementContext).unregisterMBean(with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,connector=networkConnectors,networkConnectorName=NC"))));
-            allowing(managementContext).unregisterMBean(with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,service=Log4JConfiguration"))));
-            allowing(managementContext).unregisterMBean(with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,destinationType=Topic,destinationName=ActiveMQ.Advisory.Connection"))));
-            allowing(managementContext).unregisterMBean(with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,destinationType=Topic,destinationName=ActiveMQ.Advisory.NetworkBridge"))));
-            allowing(managementContext).unregisterMBean(with(equal(
-                    new ObjectName("Test:type=Broker,brokerName=BrokerNC,destinationType=Topic,destinationName=ActiveMQ.Advisory.MasterBroker"))));
-        }});
+        managementContext = mock(ManagementContext.class);
+
+        when(managementContext.getJmxDomainName()).thenReturn("Test");
+        when(managementContext.isCreateConnector()).thenReturn(false);
+        when(managementContext.isConnectorStarted()).thenReturn(false);
+        
+        // Mock MBean registration to release semaphore - only for network connector
+        doAnswer(invocation -> {
+            ObjectName name = invocation.getArgument(1);
+            LOG.info("MBean registered: " + name);
+            if (name.toString().contains("networkConnectors")) {
+                mbeanRegistered.release();
+            }
+            return mock(ObjectInstance.class);
+        }).when(managementContext).registerMBean(any(), any(ObjectName.class));
+        
+        // Mock MBean unregistration to release semaphore - only for network connector
+        doAnswer(invocation -> {
+            ObjectName name = invocation.getArgument(0);
+            LOG.info("MBean unregistered: " + name);
+            if (name.toString().contains("networkConnectors")) {
+                mbeanUnregistered.release();
+            }
+            return null;
+        }).when(managementContext).unregisterMBean(any(ObjectName.class));
 
         brokerB = new BrokerService();
         brokerB.setManagementContext(managementContext);
@@ -177,11 +101,20 @@ public Object invoke(Invocation invocation) throws Throwable {
 
     @After
     public void tearDown() throws Exception {
-        brokerA.stop();
-        brokerA.waitUntilStopped();
-        brokerB.stop();
-        brokerB.waitUntilStopped();
-        proxy.close();
+        if (agent != null) {
+            agent.stop();
+        }
+        if (brokerB != null) {
+            brokerB.stop();
+            brokerB.waitUntilStopped();
+        }
+        if (brokerA != null) {
+            brokerA.stop();
+            brokerA.waitUntilStopped();
+        }
+        if (proxy != null) {
+            proxy.close();
+        }
     }
 
     private void configure(BrokerService broker) {
@@ -191,16 +124,18 @@ private void configure(BrokerService broker) {
 
     @Test
     public void testMulicastReconnect() throws Exception {
+        // Start the discovery agent BEFORE starting brokerB to ensure it discovers via proxy
+        agent = MulticastDiscoveryAgentFactory.createDiscoveryAgent(new URI(discoveryAddress));
+        agent.registerService(proxy.getUrl().toString());
+        agent.start();
+        
+        // Give the agent time to start advertising
+        Thread.sleep(500);
 
         brokerB.addNetworkConnector(discoveryAddress + "&discovered.trace=true&discovered.wireFormat.maxInactivityDuration=1000&discovered.wireFormat.maxInactivityDurationInitalDelay=1000");
         brokerB.start();
         brokerB.waitUntilStarted();
 
-        // control multicast advertise agent to inject proxy
-        agent = MulticastDiscoveryAgentFactory.createDiscoveryAgent(new URI(discoveryAddress));
-        agent.registerService(proxy.getUrl().toString());
-        agent.start();
-
         doReconnect();
     }
 
@@ -214,8 +149,13 @@ public void testSimpleReconnect() throws Exception {
     }
 
     private void doReconnect() throws Exception {
+        // Wait for initial network connector mbean registration
+        assertTrue("network connector mbean registered within 1 minute", 
+            mbeanRegistered.tryAcquire(60, TimeUnit.SECONDS));
 
         for (int i=0; i<maxReconnects; i++) {
+            LOG.info("Reconnect iteration: " + (i + 1));
+            
             // Wait for connection
             assertTrue("we got a network connection in a timely manner", Wait.waitFor(new Wait.Condition() {
                 @Override
@@ -224,20 +164,38 @@ public boolean isSatisified() throws Exception {
                 }
             }));
 
-            // wait for network connector
-            assertTrue("network connector mbean registered within 1 minute", mbeanRegistered.tryAcquire(60, TimeUnit.SECONDS));
-
             // force an inactivity timeout via the proxy
             proxy.pause();
-
-            // wait for the inactivity timeout and network shutdown
-            assertTrue("network connector mbean unregistered within 1 minute", mbeanUnregistered.tryAcquire(60, TimeUnit.SECONDS));
-
-            // whack all connections
+            
+            // Wait a bit longer than the inactivity duration to ensure timeout triggers
+            // maxInactivityDuration=1000ms, so wait at least 2-3x that
+            Thread.sleep(3000);
+            
+            // Now actually close the connections to force disconnect
             proxy.close();
+            
+            // Wait for connections to be cleared
+            assertTrue("proxy connections cleared", Wait.waitFor(new Wait.Condition() {
+                @Override
+                public boolean isSatisified() throws Exception {
+                   return proxy.connections.size() == 0;
+                }
+            }));
 
-            // let a reconnect succeed
-            proxy.reopen();
+            // let a reconnect succeed for next iteration
+            if (i < maxReconnects - 1) {
+                proxy.reopen();
+                // Give some time for reconnection to be attempted
+                Thread.sleep(1500);
+            }
         }
+        
+        // Now stop the broker which should unregister the network connector mbean
+        brokerB.stop();
+        brokerB.waitUntilStopped();
+        
+        // wait for the network connector mbean to be unregistered
+        assertTrue("network connector mbean unregistered within 1 minute", 
+            mbeanUnregistered.tryAcquire(60, TimeUnit.SECONDS));
     }
 }