lib/evrone/ci/worker/middlewares/docker_script.rb in evrone-ci-worker-0.2.0.pre4 vs lib/evrone/ci/worker/middlewares/docker_script.rb in evrone-ci-worker-0.2.0.pre5

@@ -1,18 +1,20 @@
-require 'net/scp'
-
 module Evrone
   module CI
     module Worker
 
       DockerScript = Struct.new(:app) do
 
         include Helper::Logger
+        include Helper::Config
+        include Common::Helper::UploadShCommand
 
         def call(env)
           if env.ssh && env.docker_repo_dir
             code = run_script(env)
+            run_after_script(env)
+
             if code == 0
               app.call env
             else
               case read_state(env)
               when "script"
@@ -27,27 +29,49 @@
         end
 
         private
 
           def run_script(env)
-            scp = ::Net::SCP.new(env.ssh.connection)
-            script = [env.docker_repo_dir, ".ci_build.sh"].join("/")
-            scp.upload! StringIO.new(content(env)), script
-            env.ssh.spawn "env - HOME=$HOME bash #{script}", chdir: env.docker_repo_dir, &env.job.method(:add_to_output)
+            home = config.docker.remote_dir
+            file = [home, ".ci_build.sh"].join("/")
+
+            script = [upload_sh_command(file, script_content(env, home))]
+            script << "env - HOME=$HOME bash #{file}"
+            script = script.join(" && ")
+
+            env.ssh.spawn script, chdir: home, &env.job.method(:add_to_output)
           end
 
-          def content(env)
+          def run_after_script(env)
+            home = config.docker.remote_dir
+            file = [home, ".ci_after_build.sh"].join("/")
+
+            script = [upload_sh_command(file, after_script_content(env))]
+            script << "env - HOME=$HOME bash #{file}"
+            script = script.join(" && ")
+
+            env.ssh.spawn script, chdir: home, &env.job.method(:add_to_output)
+          end
+
+          def script_content(env, home)
             buf = ["set -e"]
-            buf << "echo before_script > #{env.docker_repo_dir}/.ci_state"
+            buf << "echo before_script > #{home}/.ci_state"
             buf << env.job.message.before_script
-            buf << "echo script > #{env.docker_repo_dir}/.ci_state"
+            buf << "echo script > #{home}/.ci_state"
             buf << env.job.message.script
             buf.join("\n")
           end
 
+          def after_script_content(env)
+            buf = ["set -e"]
+            buf << env.job.message.after_script
+            buf.join("\n")
+          end
+
           def read_state(env)
+            home = config.docker.remote_dir
             buf = ""
-            state_file = env.docker_repo_dir.join(".ci_state")
+            state_file = "#{home}/.ci_state"
             env.ssh.spawn "cat #{state_file}" do |out|
               buf << out
             end
             buf.strip
           end