Examples with StubbableTransport org.opensearch.test.transport.StubbableTransport used on opensource projects

Example 1 with StubbableTransport

use of org.opensearch.test.transport.StubbableTransport in project OpenSearch by opensearch-project.

the class SniffConnectionStrategyTests method testConfiguredProxyAddressModeWillReplaceNodeAddress.

public void testConfiguredProxyAddressModeWillReplaceNodeAddress() {
    List<DiscoveryNode> knownNodes = new CopyOnWriteArrayList<>();
    try (MockTransportService accessible = startTransport("seed_node", knownNodes, Version.CURRENT);
        MockTransportService unresponsive1 = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool);
        MockTransportService unresponsive2 = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool)) {
        // We start in order to get a valid address + port, but do not start accepting connections as we
        // will not actually connect to these transports
        unresponsive1.start();
        unresponsive2.start();
        DiscoveryNode accessibleNode = accessible.getLocalNode();
        DiscoveryNode discoverableNode = unresponsive2.getLocalNode();
        // Use the address for the node that will not respond
        DiscoveryNode unaddressableSeedNode = new DiscoveryNode(accessibleNode.getName(), accessibleNode.getId(), accessibleNode.getEphemeralId(), accessibleNode.getHostName(), accessibleNode.getHostAddress(), unresponsive1.getLocalNode().getAddress(), accessibleNode.getAttributes(), accessibleNode.getRoles(), accessibleNode.getVersion());
        knownNodes.add(unaddressableSeedNode);
        knownNodes.add(discoverableNode);
        Collections.shuffle(knownNodes, random());
        try (MockTransportService localService = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool)) {
            localService.start();
            localService.acceptIncomingRequests();
            StubbableTransport transport = new StubbableTransport(localService.transport);
            AtomicReference<TransportAddress> discoverableNodeAddress = new AtomicReference<>();
            transport.setDefaultConnectBehavior((delegate, node, profile, listener) -> {
                if (node.equals(discoverableNode)) {
                    // Do not actually try to connect because the node will not respond. Just capture the
                    // address for later assertion
                    discoverableNodeAddress.set(node.getAddress());
                    listener.onFailure(new ConnectTransportException(node, "general failure"));
                } else {
                    delegate.openConnection(node, profile, listener);
                }
            });
            List<String> seedNodes = Collections.singletonList(accessibleNode.toString());
            TransportAddress proxyAddress = accessibleNode.getAddress();
            ClusterConnectionManager connectionManager = new ClusterConnectionManager(profile, transport);
            try (RemoteConnectionManager remoteConnectionManager = new RemoteConnectionManager(clusterAlias, connectionManager);
                SniffConnectionStrategy strategy = new SniffConnectionStrategy(clusterAlias, localService, remoteConnectionManager, proxyAddress.toString(), Settings.EMPTY, 3, n -> true, seedNodes)) {
                assertFalse(connectionManager.nodeConnected(unaddressableSeedNode));
                assertFalse(connectionManager.nodeConnected(discoverableNode));
                assertTrue(strategy.assertNoRunningConnections());
                PlainActionFuture<Void> connectFuture = PlainActionFuture.newFuture();
                strategy.connect(connectFuture);
                connectFuture.actionGet();
                assertTrue(connectionManager.nodeConnected(unaddressableSeedNode));
                // Connection to discoverable will fail due to the stubbable transport
                assertFalse(connectionManager.nodeConnected(discoverableNode));
                assertEquals(proxyAddress, discoverableNodeAddress.get());
                assertTrue(strategy.assertNoRunningConnections());
            }
        }
    }
}