Merge pull request #3983 from fviale/master

Fix for Node source permissions lost after scheduler kill

rm/rm-server/src/main/java/org/ow2/proactive/resourcemanager/core/RMCore.java

@@ -1628,6 +1628,7 @@ public class RMCore implements ResourceManager, InitActive, RunActive {
         NodeSource nodeSource = this.createNodeSourceInstance(nodeSourceDescriptor);
 
         this.definedNodeSources.put(nodeSourceDescriptor.getName(), nodeSource);
+        this.emitNodeSourceEvent(nodeSource, RMEventType.NODESOURCE_DEFINED);
 
         if (nodeSourceDescriptor.getStatus().equals(NodeSourceStatus.NODES_DEPLOYED)) {
             NodeSource nodeSourceToDeploy = this.createNodeSourceInstance(nodeSourceDescriptor);
@@ -1657,6 +1658,7 @@ public class RMCore implements ResourceManager, InitActive, RunActive {
                                              nodeSourcePolicyStub);
 
             this.deployedNodeSources.put(nodeSourceName, nodeSourceStub);
+            this.emitNodeSourceEvent(nodeSourceStub, RMEventType.NODESOURCE_CREATED);
         } else {
             this.nodesRecoveryManager.logRecoveryAbortedReason(nodeSourceName, "This node source is undeployed");
         }
@@ -2011,7 +2013,14 @@ public class RMCore implements ResourceManager, InitActive, RunActive {
 
         final List<RMNodeSourceEvent> nodeSourceEvents = new ArrayList<>(this.definedNodeSources.values()
                                                                                                 .stream()
-                                                                                                .map(NodeSource::createNodeSourceEvent)
+                                                                                                .map(ns -> {
+                                                                                                    if (deployedNodeSources.containsKey(ns.getName())) {
+                                                                                                        return PAFuture.getFutureValue(deployedNodeSources.get(ns.getName())
+                                                                                                                                                          .createNodeSourceEvent());
+                                                                                                    } else {
+                                                                                                        return ns.createNodeSourceEvent();
+                                                                                                    }
+                                                                                                })
                                                                                                 .collect(Collectors.toList()));
 
         long eventCounter = 0;

rm/rm-server/src/main/java/org/ow2/proactive/resourcemanager/core/recovery/NodesRecoveryManager.java

@@ -51,11 +51,13 @@ import org.ow2.proactive.resourcemanager.core.RMCore;
 import org.ow2.proactive.resourcemanager.core.properties.PAResourceManagerProperties;
 import org.ow2.proactive.resourcemanager.db.NodeSourceData;
 import org.ow2.proactive.resourcemanager.db.RMNodeData;
+import org.ow2.proactive.resourcemanager.exception.AddingNodesException;
 import org.ow2.proactive.resourcemanager.nodesource.NodeSource;
 import org.ow2.proactive.resourcemanager.nodesource.NodeSourceDescriptor;
 import org.ow2.proactive.resourcemanager.nodesource.infrastructure.InfrastructureManager;
 import org.ow2.proactive.resourcemanager.nodesource.infrastructure.InfrastructureManagerFactory;
 import org.ow2.proactive.resourcemanager.rmnode.RMNode;
+import org.ow2.proactive.resourcemanager.utils.RMNodeStarter;
 import org.ow2.proactive.utils.PAExecutors;
 
 import com.google.common.util.concurrent.Uninterruptibles;
@@ -227,8 +229,22 @@ public class NodesRecoveryManager {
     private synchronized RMNode addRMNodeToCoreAndSource(NodeSource nodeSource, Map<NodeState, Integer> nodeStates,
             RMNodeData rmNodeData, String nodeUrl, Node node, NodeState previousState) {
         RMNode rmNode = nodeSource.internalAddNodeAfterRecovery(node, rmNodeData);
+        boolean tokenInNodeSource = nodeSource.getNodeUserAccessType().getTokens() != null &&
+                                    nodeSource.getNodeUserAccessType().getTokens().length > 0;
+        boolean tokenInNode = false;
         this.rmCore.registerAvailableNode(rmNode);
         if (!(node instanceof FakeDownNodeForRecovery)) {
+            try {
+                String nodeAccessToken = node.getProperty(RMNodeStarter.NODE_ACCESS_TOKEN);
+                tokenInNode = nodeAccessToken != null && nodeAccessToken.length() > 0;
+
+                if (tokenInNode) {
+                    logger.debug("Node " + node.getNodeInformation().getURL() + " is protected by access token " +
+                                 nodeAccessToken);
+                }
+            } catch (Exception e) {
+                throw new AddingNodesException(e);
+            }
             try {
                 RMCore.topologyManager.addNode(rmNode.getNode());
             } catch (Exception e) {
@@ -240,6 +256,7 @@ public class NodesRecoveryManager {
             logger.info("Triggering down node notification for " + nodeUrl);
             this.triggerDownNodeHookIfNecessary(nodeSource, rmNodeData, nodeUrl, previousState);
         }
+        rmNode.setProtectedByToken(tokenInNode || tokenInNodeSource);
         this.updateRecoveredNodeStateCounter(nodeStates, rmNode.getState());
         return rmNode;
     }

rm/rm-server/src/main/java/org/ow2/proactive/resourcemanager/nodesource/NodeSource.java

@@ -250,6 +250,10 @@ public class NodeSource implements InitActive, RunActive {
                                              .orElse(new LinkedHashMap<>());
     }
 
+    public AccessType getNodeUserAccessType() {
+        return nodeUserAccessType;
+    }
+
     public static void initThreadPools() {
         if (threadPoolHolder == null) {
             try {

rm/rm-server/src/main/java/org/ow2/proactive/resourcemanager/rmnode/RMDeployingNode.java

@@ -372,6 +372,11 @@ public final class RMDeployingNode extends AbstractRMNode {
         return false;
     }
 
+    @Override
+    public void setProtectedByToken(boolean protectedByToken) {
+
+    }
+
     public RMDeployingNode updateOnNodeSource() {
         return nodeSource.update(this);
     }

rm/rm-server/src/main/java/org/ow2/proactive/resourcemanager/rmnode/RMNode.java

@@ -352,6 +352,8 @@ public interface RMNode extends Comparable<RMNode> {
      */
     boolean isProtectedByToken();
 
+    void setProtectedByToken(boolean protectedByToken);
+
     void addToken(String token);
 
     void removeToken(String token);

rm/rm-server/src/main/java/org/ow2/proactive/resourcemanager/rmnode/RMNodeImpl.java

@@ -489,6 +489,7 @@ public class RMNodeImpl extends AbstractRMNode {
         return protectedByToken;
     }
 
+    @Override
     public void setProtectedByToken(boolean protectedByToken) {
         this.protectedByToken = protectedByToken;
     }

rm/rm-server/src/test/java/functionaltests/RegressionTestSuite.java

@@ -34,6 +34,7 @@ import functionaltests.nodesource.TestLocalInfrastructureTimeSlotPolicy;
 import functionaltests.nodesource.TestNodeSourceThreadPool;
 import functionaltests.nodesource.TestSSHInfrastructureV2RestartDownNodesPolicy;
 import functionaltests.nodesrecovery.NodesRecoveryPropertyTest;
+import functionaltests.nodesrecovery.RecoverDefaultInfrastructureTest;
 import functionaltests.nodesrecovery.RecoverLocalInfrastructureTest;
 import functionaltests.nodesrecovery.RecoverSSHInfrastructureV2Test;
 import functionaltests.nonblockingcore.NonBlockingCoreTest;
@@ -44,9 +45,10 @@ import functionaltests.selectionscript.SelectionScriptTimeOutTest;
 @RunWith(Suite.class)
 @Suite.SuiteClasses({ NonBlockingCoreTest.class, TestNSNodesPermissions.class, SelectionScriptTimeOutTest.class,
                       RecoverLocalInfrastructureTest.class, RecoverSSHInfrastructureV2Test.class,
-                      TestLocalInfrastructureCronPolicy.class, TestLocalInfrastructureTimeSlotPolicy.class,
-                      TestSSHInfrastructureV2RestartDownNodesPolicy.class, NodesRecoveryPropertyTest.class,
-                      NodesHouseKeepingPropertiesTest.class, TestNodeSourceThreadPool.class })
+                      RecoverDefaultInfrastructureTest.class, TestLocalInfrastructureCronPolicy.class,
+                      TestLocalInfrastructureTimeSlotPolicy.class, TestSSHInfrastructureV2RestartDownNodesPolicy.class,
+                      NodesRecoveryPropertyTest.class, NodesHouseKeepingPropertiesTest.class,
+                      TestNodeSourceThreadPool.class })
 
 /**
  * @author ActiveEon Team

rm/rm-server/src/test/java/functionaltests/nodesrecovery/RecoverDefaultInfrastructureTest.java

+/*
+ * ProActive Parallel Suite(TM):
+ * The Open Source library for parallel and distributed
+ * Workflows & Scheduling, Orchestration, Cloud Automation
+ * and Big Data Analysis on Enterprise Grids & Clouds.
+ *
+ * Copyright (c) 2007 - 2017 ActiveEon
+ * Contact: contact@activeeon.com
+ *
+ * This library is free software: you can redistribute it and/or
+ * modify it under the terms of the GNU Affero General Public License
+ * as published by the Free Software Foundation: version 3 of
+ * the License.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ *
+ * If needed, contact us to obtain a release under GPL Version 2 or 3
+ * or a different license than the AGPL.
+ */
+package functionaltests.nodesrecovery;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import java.io.File;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.objectweb.proactive.core.util.wrapper.BooleanWrapper;
+import org.ow2.proactive.resourcemanager.common.NodeState;
+import org.ow2.proactive.resourcemanager.common.event.RMEventType;
+import org.ow2.proactive.resourcemanager.common.event.RMNodeEvent;
+import org.ow2.proactive.resourcemanager.common.event.RMNodeSourceEvent;
+import org.ow2.proactive.resourcemanager.core.properties.PAResourceManagerProperties;
+import org.ow2.proactive.resourcemanager.frontend.ResourceManager;
+import org.ow2.proactive.resourcemanager.nodesource.NodeSourceStatus;
+import org.ow2.proactive.utils.Criteria;
+import org.ow2.proactive.utils.NodeSet;
+
+import functionaltests.monitor.RMMonitorEventReceiver;
+import functionaltests.utils.RMFunctionalTest;
+import functionaltests.utils.RMTHelper;
+import functionaltests.utils.TestNode;
+
+
+/**
+ * @author ActiveEon Team
+ * @since 22/06/17
+ */
+public class RecoverDefaultInfrastructureTest extends RMFunctionalTest {
+
+    private static final String START_CONFIG = "/functionaltests/config/functionalTRMProperties-RM-start-clean-db-nodes-recovery-enabled.ini";
+
+    private static final String RESTART_CONFIG = "/functionaltests/config/functionalTRMProperties-RM-restart-keep-db-nodes-recovery-enabled.ini";
+
+    private static final String NODE_SOURCE_NAME = "Default_" + RecoverDefaultInfrastructureTest.class.getSimpleName();
+
+    private static final String TOKEN = "toto";
+
+    private static final String NODE_SOURCE_USER_ACCESS_TYPE = "tokens=" + TOKEN;
+
+    private static final int NODE_NUMBER = 3;
+
+    private ResourceManager resourceManager = null;
+
+    @Before
+    public void setup() throws Exception {
+
+        this.startRmAndCheckInitialState();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        try {
+            RecoverInfrastructureTestHelper.killNodesWithStrongSigKill();
+        } catch (NodesRecoveryProcessHelper.ProcessNotFoundException e) {
+            RMTHelper.log("Cannot kill the node process: " + e.getMessage());
+        }
+    }
+
+    @Test
+    public void testRecoverDefaultInfrastructureWithAliveNodes() throws Exception {
+
+        this.createNodeSourceAndCheckState();
+
+        getNodesFromRMWithToken();
+
+        RecoverInfrastructureTestHelper.killRmWithStrongSigKill();
+
+        this.restartRmAndCheckFinalState();
+
+        this.checkNodesStateAfterRecovery(NODE_NUMBER, 0);
+
+        getNodesFromRMWithToken();
+    }
+
+    @Test
+    public void testRecoverUndeployedNodeSource() throws Exception {
+
+        this.defineNodeSourceAndCheckState();
+
+        this.restartRmAndCheckFinalState();
+
+        this.checkNodesStateAfterRecovery(0, 0);
+    }
+
+    private void getNodesFromRMWithToken() throws Exception {
+        Criteria criteria = new Criteria(NODE_NUMBER);
+        criteria.setNodeAccessToken(TOKEN);
+
+        RMTHelper.log("Acquiring " + NODE_NUMBER + " nodes with token=" + TOKEN);
+        NodeSet nodeSet = this.rmHelper.getResourceManager().getNodes(criteria);
+        Assert.assertEquals(NODE_NUMBER, nodeSet.size());
+        this.rmHelper.getResourceManager().releaseNodes(nodeSet);
+        this.rmHelper.waitForAnyMultipleNodeEvent(RMEventType.NODE_STATE_CHANGED, NODE_NUMBER);
+    }
+
+    private void startRmWithConfig(String configurationFilePath) throws Exception {
+        String rmconf = new File(RMTHelper.class.getResource(configurationFilePath).toURI()).getAbsolutePath();
+        this.rmHelper.startRM(rmconf);
+        this.resourceManager = this.rmHelper.getResourceManager();
+    }
+
+    private void startRmAndCheckInitialState() throws Exception {
+        // start RM
+        this.startRmWithConfig(START_CONFIG);
+        assertThat(PAResourceManagerProperties.RM_PRESERVE_NODES_ON_SHUTDOWN.getValueAsBoolean()).isTrue();
+        assertThat(this.rmHelper.isRMStarted()).isTrue();
+
+        // check the initial state of the RM
+        assertThat(this.resourceManager.getState().getAllNodes().size()).isEqualTo(0);
+    }
+
+    private void createNodeSourceAndCheckState() throws Exception {
+        this.rmHelper.createDefaultNodeSourceWithNodesRecoverable(NODE_SOURCE_NAME,
+                                                                  new Object[] { NODE_SOURCE_USER_ACCESS_TYPE, "ME" });
+        RMMonitorEventReceiver resourceManagerMonitor = (RMMonitorEventReceiver) this.resourceManager;
+
+        testNodes.addAll(this.rmHelper.addNodesToDefaultNodeSource(NODE_SOURCE_NAME, NODE_NUMBER));
+
+        List<RMNodeSourceEvent> nodeSourceEventPerNodeSource = resourceManagerMonitor.getInitialState()
+                                                                                     .getNodeSourceEvents();
+        assertThat(nodeSourceEventPerNodeSource.size()).isEqualTo(1);
+
+        RMNodeSourceEvent rmNodeSourceEvent = nodeSourceEventPerNodeSource.get(0);
+        assertThat(rmNodeSourceEvent.getSourceName()).isEqualTo(NODE_SOURCE_NAME);
+        assertThat(rmNodeSourceEvent.getNodeSourceStatus()).isEqualTo(NodeSourceStatus.NODES_DEPLOYED.toString());
+
+        assertThat(resourceManagerMonitor.getState().getAllNodes().size()).isEqualTo(NODE_NUMBER);
+    }
+
+    private void defineNodeSourceAndCheckState() throws Exception {
+        this.rmHelper.defineNodeSource(NODE_SOURCE_NAME, NODE_NUMBER);
+        RMMonitorEventReceiver resourceManagerMonitor = (RMMonitorEventReceiver) this.resourceManager;
+        List<RMNodeSourceEvent> nodeSourceEventPerNodeSource = resourceManagerMonitor.getInitialState()
+                                                                                     .getNodeSourceEvents();
+        assertThat(nodeSourceEventPerNodeSource.size()).isEqualTo(1);
+        assertThat(nodeSourceEventPerNodeSource.get(0).getSourceName()).isEqualTo(NODE_SOURCE_NAME);
+        assertThat(nodeSourceEventPerNodeSource.get(0)
+                                               .getNodeSourceDescription()
+                                               .contains(NODE_SOURCE_USER_ACCESS_TYPE));
+    }
+
+    private void restartRmAndCheckFinalState() throws Exception {
+        // restart RM
+        this.rmHelper = new RMTHelper();
+        this.startRmWithConfig(RESTART_CONFIG);
+        assertThat(PAResourceManagerProperties.RM_PRESERVE_NODES_ON_SHUTDOWN.getValueAsBoolean()).isFalse();
+        assertThat(this.rmHelper.isRMStarted()).isTrue();
+
+        // re-snapshot the RM state
+        RMMonitorEventReceiver resourceManagerMonitor = (RMMonitorEventReceiver) this.resourceManager;
+        List<RMNodeSourceEvent> nodeSourceEvent = resourceManagerMonitor.getInitialState().getNodeSourceEvents();
+
+        // the node source has been recovered on restart: we should have one node source with the same name
+        assertThat(nodeSourceEvent.size()).isEqualTo(1);
+        assertThat(nodeSourceEvent.get(0).getSourceName()).isEqualTo(NODE_SOURCE_NAME);
+        assertThat(nodeSourceEvent.get(0).getNodeSourceDescription().contains(NODE_SOURCE_USER_ACCESS_TYPE));
+    }
+
+    private void checkNodesStateAfterRecovery(int expectedNbAliveNodes, int expectedNbDownNodes) {
+        RMMonitorEventReceiver resourceManagerMonitor = (RMMonitorEventReceiver) this.resourceManager;
+        // the nodes should have been recovered too, and should be alive
+        Set<String> allNodes = resourceManagerMonitor.getState().getAllNodes();
+        assertThat(allNodes.size()).isEqualTo(expectedNbAliveNodes + expectedNbDownNodes);
+        Set<String> nodeSourceNames = new HashSet<>();
+        nodeSourceNames.add(NODE_SOURCE_NAME);
+        Set<String> aliveNodeUrls = this.resourceManager.listAliveNodeUrls(nodeSourceNames);
+        assertThat(aliveNodeUrls.size()).isEqualTo(expectedNbAliveNodes);
+
+        List<RMNodeEvent> nodeEvents = resourceManagerMonitor.getInitialState().getNodeEvents();
+        long nbFreeNodes = nodeEvents.stream()
+                                     .filter(nodeEvent -> nodeEvent.getNodeState().equals(NodeState.FREE))
+                                     .count();
+        assertThat(nbFreeNodes).isEqualTo(expectedNbAliveNodes);
+        long nbDownNodes = nodeEvents.stream()
+                                     .filter(nodeEvent -> nodeEvent.getNodeState().equals(NodeState.DOWN))
+                                     .count();
+        assertThat(nbDownNodes).isEqualTo(expectedNbDownNodes);
+
+        if (expectedNbAliveNodes > 0) {
+            // the recovered nodes should be usable, try to lock/unlock them to see
+            BooleanWrapper lockSucceeded = this.resourceManager.lockNodes(allNodes);
+            assertThat(lockSucceeded).isEqualTo(new BooleanWrapper(true));
+            BooleanWrapper unlockSucceeded = this.resourceManager.unlockNodes(allNodes);
+            assertThat(unlockSucceeded).isEqualTo(new BooleanWrapper(true));
+        }
+    }
+
+}

rm/rm-server/src/test/java/functionaltests/nodesrecovery/RecoverLocalInfrastructureTest.java

@@ -33,6 +33,7 @@ import java.util.List;
 import java.util.Set;
 
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.objectweb.proactive.core.util.wrapper.BooleanWrapper;
@@ -43,6 +44,8 @@ import org.ow2.proactive.resourcemanager.common.event.RMNodeSourceEvent;
 import org.ow2.proactive.resourcemanager.core.properties.PAResourceManagerProperties;
 import org.ow2.proactive.resourcemanager.frontend.ResourceManager;
 import org.ow2.proactive.resourcemanager.nodesource.NodeSourceStatus;
+import org.ow2.proactive.utils.Criteria;
+import org.ow2.proactive.utils.NodeSet;
 
 import functionaltests.monitor.RMMonitorEventReceiver;
 import functionaltests.utils.RMFunctionalTest;
@@ -62,6 +65,10 @@ public class RecoverLocalInfrastructureTest extends RMFunctionalTest {
     private static final String NODE_SOURCE_NAME = "LocalNodeSource" +
                                                    RecoverLocalInfrastructureTest.class.getSimpleName();
 
+    private static final String TOKEN = "toto";
+
+    private static final String NODE_SOURCE_USER_ACCESS_TYPE = "tokens=" + TOKEN;
+
     private static final int NODE_NUMBER = 3;
 
     private ResourceManager resourceManager = null;
@@ -86,11 +93,15 @@ public class RecoverLocalInfrastructureTest extends RMFunctionalTest {
 
         this.createNodeSourceAndCheckState();
 
+        getNodesFromRMWithToken();
+
         RecoverInfrastructureTestHelper.killRmWithStrongSigKill();
 
         this.restartRmAndCheckFinalState();
 
         this.checkNodesStateAfterRecovery(NODE_NUMBER, 0);
+
+        getNodesFromRMWithToken();
     }
 
     @Test
@@ -132,7 +143,9 @@ public class RecoverLocalInfrastructureTest extends RMFunctionalTest {
     }
 
     private void createNodeSourceAndCheckState() throws Exception {
-        this.rmHelper.createNodeSourceWithNodesRecoverable(NODE_SOURCE_NAME, NODE_NUMBER);
+        this.rmHelper.createNodeSourceWithNodesRecoverable(NODE_SOURCE_NAME,
+                                                           NODE_NUMBER,
+                                                           new Object[] { "tokens=toto", "ME" });
         RMMonitorEventReceiver resourceManagerMonitor = (RMMonitorEventReceiver) this.resourceManager;
 
         List<RMNodeSourceEvent> nodeSourceEventPerNodeSource = resourceManagerMonitor.getInitialState()
@@ -153,6 +166,20 @@ public class RecoverLocalInfrastructureTest extends RMFunctionalTest {
                                                                                      .getNodeSourceEvents();
         assertThat(nodeSourceEventPerNodeSource.size()).isEqualTo(1);
         assertThat(nodeSourceEventPerNodeSource.get(0).getSourceName()).isEqualTo(NODE_SOURCE_NAME);
+        assertThat(nodeSourceEventPerNodeSource.get(0)
+                                               .getNodeSourceDescription()
+                                               .contains(NODE_SOURCE_USER_ACCESS_TYPE));
+    }
+
+    private void getNodesFromRMWithToken() throws Exception {
+        Criteria criteria = new Criteria(NODE_NUMBER);
+        criteria.setNodeAccessToken(TOKEN);
+
+        RMTHelper.log("Acquiring " + NODE_NUMBER + " nodes with token=" + TOKEN);
+        NodeSet nodeSet = this.rmHelper.getResourceManager().getNodes(criteria);
+        Assert.assertEquals(NODE_NUMBER, nodeSet.size());
+        this.rmHelper.getResourceManager().releaseNodes(nodeSet);
+        this.rmHelper.waitForAnyMultipleNodeEvent(RMEventType.NODE_STATE_CHANGED, NODE_NUMBER);
     }
 
     private void restartRmAndCheckFinalState() throws Exception {
@@ -169,6 +196,7 @@ public class RecoverLocalInfrastructureTest extends RMFunctionalTest {
         // the node source has been recovered on restart: we should have one node source with the same name
         assertThat(nodeSourceEvent.size()).isEqualTo(1);
         assertThat(nodeSourceEvent.get(0).getSourceName()).isEqualTo(NODE_SOURCE_NAME);
+        assertThat(nodeSourceEvent.get(0).getNodeSourceDescription().contains(NODE_SOURCE_USER_ACCESS_TYPE));
     }
 
     private void checkNodesStateAfterRecovery(int expectedNbAliveNodes, int expectedNbDownNodes) {

rm/rm-server/src/test/java/functionaltests/utils/RMTHelper.java

@@ -56,6 +56,7 @@ import org.ow2.proactive.resourcemanager.common.event.RMNodeEvent;
 import org.ow2.proactive.resourcemanager.core.properties.PAResourceManagerProperties;
 import org.ow2.proactive.resourcemanager.frontend.ResourceManager;
 import org.ow2.proactive.resourcemanager.nodesource.NodeSource;
+import org.ow2.proactive.resourcemanager.nodesource.infrastructure.DefaultInfrastructureManager;
 import org.ow2.proactive.resourcemanager.nodesource.infrastructure.LocalInfrastructure;
 import org.ow2.proactive.resourcemanager.nodesource.policy.StaticPolicy;
 import org.ow2.proactive.utils.FileToBytesConverter;
@@ -147,7 +148,20 @@ public class RMTHelper {
     }
 
     public void createNodeSourceWithNodesRecoverable(String name, int nodeNumber) throws Exception {
-        createNodeSourceWithNodesRecoverable(name, nodeNumber, getResourceManager(), getMonitorsHandler());
+        createNodeSourceWithNodesRecoverable(name, nodeNumber, null);
+    }
+
+    public void createNodeSourceWithNodesRecoverable(String name, int nodeNumber, Object[] policyParameters)
+            throws Exception {
+        createNodeSourceWithNodesRecoverable(name,
+                                             nodeNumber,
+                                             getResourceManager(),
+                                             getMonitorsHandler(),
+                                             policyParameters);
+    }
+
+    public void createDefaultNodeSourceWithNodesRecoverable(String name, Object[] policyParameters) throws Exception {
+        createDefaultNodeSourceWithNodesRecoverable(name, getResourceManager(), getMonitorsHandler(), policyParameters);
     }
 
     public void createNodeSourceWithInfiniteTimeout(String name, int nodeNumber) throws Exception {
@@ -156,11 +170,11 @@ public class RMTHelper {
 
     public static void createNodeSource(String name, int nodeNumber, List<String> vmOptions, ResourceManager rm,
             RMMonitorsHandler monitor) throws Exception {
-        createNodeSource(name, nodeNumber, vmOptions, rm, monitor, NODES_NOT_RECOVERABLE);
+        createNodeSource(name, nodeNumber, vmOptions, rm, monitor, NODES_NOT_RECOVERABLE, null);
     }
 
     public static void createNodeSource(String name, int nodeNumber, List<String> vmOptions, ResourceManager rm,
-            RMMonitorsHandler monitor, boolean nodesRecoverable) throws Exception {
+            RMMonitorsHandler monitor, boolean nodesRecoverable, Object[] policyParameters) throws Exception {
         byte[] creds = setJavaPropertyAndGetCredentials();
         log("Creating a node source " + name);
         rm.createNodeSource(name,
@@ -169,13 +183,28 @@ public class RMTHelper {
                                            vmOptions != null ? setup.listToString(vmOptions)
                                                              : setup.getJvmParameters() },
                             StaticPolicy.class.getName(),
-                            null,
+                            policyParameters,
                             nodesRecoverable);
         rm.setNodeSourcePingFrequency(5000, name);
 
         waitForNodeSourceCreation(name, nodeNumber, monitor);
     }
 
+    public static void createDefaultInfrastructureNodeSource(String name, ResourceManager rm, RMMonitorsHandler monitor,
+            boolean nodesRecoverable, Object[] policyParameters) throws Exception {
+        byte[] creds = setJavaPropertyAndGetCredentials();
+        log("Creating a node source " + name);
+        rm.createNodeSource(name,
+                            DefaultInfrastructureManager.class.getName(),
+                            new Object[] {},
+                            StaticPolicy.class.getName(),
+                            policyParameters,
+                            nodesRecoverable);
+        rm.setNodeSourcePingFrequency(5000, name);
+
+        waitForNodeSourceCreation(name, 0, monitor);
+    }
+
     public void defineNodeSource(String name, int nodeNumber) throws Exception {
         defineNodeSource(name,
                          nodeNumber,
@@ -232,12 +261,23 @@ public class RMTHelper {
      */
     public static void createNodeSource(String name, int nodeNumber, ResourceManager rm, RMMonitorsHandler monitor)
             throws Exception {
-        createNodeSource(name, nodeNumber, setup.getJvmParametersAsList(), rm, monitor, NODES_NOT_RECOVERABLE);
+        createNodeSource(name, nodeNumber, setup.getJvmParametersAsList(), rm, monitor, NODES_NOT_RECOVERABLE, null);
     }
 
     public static void createNodeSourceWithNodesRecoverable(String name, int nodeNumber, ResourceManager rm,
-            RMMonitorsHandler monitor) throws Exception {
-        createNodeSource(name, nodeNumber, setup.getJvmParametersAsList(), rm, monitor, NODES_RECOVERABLE);
+            RMMonitorsHandler monitor, Object[] policyParameters) throws Exception {
+        createNodeSource(name,
+                         nodeNumber,
+                         setup.getJvmParametersAsList(),
+                         rm,
+                         monitor,
+                         NODES_RECOVERABLE,
+                         policyParameters);
+    }
+
+    public static void createDefaultNodeSourceWithNodesRecoverable(String name, ResourceManager rm,
+            RMMonitorsHandler monitor, Object[] policyParameters) throws Exception {
+        createDefaultInfrastructureNodeSource(name, rm, monitor, NODES_RECOVERABLE, policyParameters);
     }
 
     public static void createNodeSourceWithInfiniteTimeout(String name, int nodeNumber, ResourceManager rm,
@@ -249,8 +289,22 @@ public class RMTHelper {
         return addNodesToDefaultNodeSource(nodesNumber, new ArrayList<String>());
     }
 
+    public List<TestNode> addNodesToDefaultNodeSource(String nodeSourceName, int nodesNumber) throws Exception {
+        return addNodesToDefaultNodeSource(nodeSourceName, nodesNumber, new ArrayList<String>());
+    }
+
     public List<TestNode> addNodesToDefaultNodeSource(int nodesNumber, List<String> vmOptions) throws Exception {
-        return addNodesToDefaultNodeSource(nodesNumber,
+        return addNodesToDefaultNodeSource(null,
+                                           nodesNumber,
+                                           vmOptions != null ? vmOptions : setup.getJvmParametersAsList(),
+                                           getResourceManager(),
+                                           getMonitorsHandler());
+    }
+
+    public List<TestNode> addNodesToDefaultNodeSource(String nodeSourceName, int nodesNumber, List<String> vmOptions)
+            throws Exception {
+        return addNodesToDefaultNodeSource(nodeSourceName,
+                                           nodesNumber,
                                            vmOptions != null ? vmOptions : setup.getJvmParametersAsList(),
                                            getResourceManager(),
                                            getMonitorsHandler());