refactor init script / user data resolution into method and add test

src/main/groovy/bpipe/executor/AWSEC2CommandExecutor.groovy

@@ -349,7 +349,6 @@ class AWSEC2CommandExecutor extends CloudExecutor {
 
         RunInstancesRequest runInstancesRequest = new RunInstancesRequest();
         runInstancesRequest.withImageId(image)
-//                .withInstanceType(InstanceType.T1Micro)
                 .withInstanceType(instanceType)
                 .withMinCount(1)
                 .withMaxCount(1)
@@ -366,14 +365,11 @@ class AWSEC2CommandExecutor extends CloudExecutor {
 
         if(config.containsKey('instanceProfile'))
             runInstancesRequest.setIamInstanceProfile(new IamInstanceProfileSpecification().withName((String)config.instanceProfile))
-
-        if(config.containsKey('initScript')) {
-            final String script = config['initScript']
-            final String userData = script.bytes.encodeBase64()
+
+        def userData = resolveUserData(config)
+        if(userData)
             runInstancesRequest = runInstancesRequest.withUserData(userData)
-            log.info "Configured init script for $command.name as userData"
-        }         
-
+
         RunInstancesResult result       
 
         List<String> retryErrorCodes = [ "InsufficientInstanceCapacity", "VcpuLimitExceeded", "RequestLimitExceeded" ]
@@ -393,6 +389,41 @@ class AWSEC2CommandExecutor extends CloudExecutor {
         log.info "Instance $instanceId started with host name $hostname"
     }
 
+    /**
+     * Check the given config for elements that would require user data to be set. These include
+     * an explicit init script, or alternatively, a walltime which is implemented through a shutdown
+     * command that is set as an init script
+     * 
+     * @param config
+     * @return  base 64 encoded user data
+     */
+    @CompileStatic
+    String resolveUserData(Map config) {
+        List<String> initScriptParts = []
+        if(config.containsKey('walltime')) {
+            // Walltime expressed in hh:mm:ss form
+            // But we need it in minutes
+            List<Integer> timeParts = ((String)config.walltime).tokenize(':')*.toInteger()
+            int minutes = timeParts.size()==3 ? (timeParts[0] * 60 + timeParts[1]) : timeParts[0]
+            log.info "Calculated $minutes shutdown minutes from walltime $config.walltime"
+            initScriptParts.add("shutdown -h +$minutes".toString())
+        }
+
+        if(config.containsKey('initScript')) {
+            initScriptParts.add(config['initScript'])
+        }
+
+        if(initScriptParts) {
+            final String script = '#!/bin/bash\n' + initScriptParts.join('\n') + '\n'
+            final String userData = script.bytes.encodeBase64()
+            log.info "Resolved init script for $command.name as userData:\n\n$script"
+            return userData
+        }         
+        else {
+            return null
+        }
+    }
+
     @CompileStatic
     void connectInstance(Map config) {
 

test-src/bpipe/executor/AWSEC2CommandExecutorTest.groovy

@@ -24,6 +24,7 @@ class AWSEC2CommandExecutorTest {
             user: 'centos',
             region: 'ap-southeast-2'
    ]
+
 
     @BeforeClass
     static void before() {
@@ -79,4 +80,37 @@ class AWSEC2CommandExecutorTest {
         println "Terminating instance $awse.instanceId"
 //        awse.ec2.terminateInstances(new TerminateInstancesRequest([awse.instanceId]))
     }
+
+
+    @Test
+    public void 'user data calculation'() {
+
+        def awse = new AWSEC2CommandExecutor(command: new Command(name: 'test'))
+
+        String userData = awse.resolveUserData(walltime:'12:00:00')
+        assert userData : "should have user data when walltime is set"
+
+        String result = new String(userData.decodeBase64(), 'utf-8')
+        assert result.contains("shutdown -h +720")
+
+        userData = awse.resolveUserData(walltime:'12:00')
+        assert userData : "should have user data when walltime is set"
+        result = new String(userData.decodeBase64(), 'utf-8')
+        assert result.contains("shutdown -h +12")
+
+        userData = awse.resolveUserData(walltime:'12:00', initScript: 'echo "hello world"')
+        assert userData : "should have user data when walltime or initScript is set"
+        result = new String(userData.decodeBase64(), 'utf-8')
+        assert result.contains("shutdown -h +12")
+
+        userData = awse.resolveUserData(initScript: 'echo "hello world"')
+        assert userData : "should have user data when walltime or initScript is set"
+        result = new String(userData.decodeBase64(), 'utf-8')
+        assert !result.contains("shutdown")
+        assert result.contains("hello world")
+
+        userData = awse.resolveUserData([:])
+        assert !userData : "should not have user data when neither walltime or initScript is set"
+
+    }
 }