Merge branch 'main' into customizable-audit-broadcast

Dockerfile

@@ -14,7 +14,7 @@ COPY Gemfile Gemfile.lock mrsk.gemspec ./
 COPY lib/mrsk/version.rb /mrsk/lib/mrsk/version.rb
 
 # Install system dependencies
-RUN apk add --no-cache --update build-base git docker openrc \
+RUN apk add --no-cache --update build-base git docker openrc openssh-client-default \
     && rc-update add docker boot \
     && gem install bundler --version=2.4.3 \
     && bundle install
@@ -31,6 +31,10 @@ RUN gem build mrsk.gemspec && \
 # Set the working directory to /workdir
 WORKDIR /workdir
 
+# Tell git it's safe to access /workdir/.git even if
+# the directory is owned by a different user
+RUN git config --global --add safe.directory /workdir
+
 # Set the entrypoint to run the installed binary in /workdir
 # Example:  docker run -it -v "$PWD:/workdir" mrsk init
 ENTRYPOINT ["mrsk"]

README.md

@@ -686,9 +686,11 @@ MRSK currently sets:
 - `MRSK_ROLE` - the specific role being targetted, if any
 - `MRSK_MESSAGE` - full text of the action (e.g. "Deployed app@150b24f")
 
-### Custom healthcheck
+### Healthcheck
 
-MRSK defaults to checking the health of your application again `/up` on port 3000 up to 7 times. You can tailor the behaviour with the `healthcheck` setting:
+MRSK uses Docker healtchecks to check the health of your application during deployment. Traefik uses this same healthcheck status to determine when a container is ready to receive traffic.
+
+The healthcheck defaults to testing the HTTP response to the path `/up` on port 3000, up to 7 times. You can tailor this behaviour with the `healthcheck` setting:
 
 ```yaml
 healthcheck:
@@ -699,7 +701,29 @@ healthcheck:
 
 This will ensure your application is configured with a traefik label for the healthcheck against `/healthz` and that the pre-deploy healthcheck that MRSK performs is done against the same path on port 4000.
 
-The healthcheck also allows for an optional `max_attempts` setting, which will attempt the healthcheck up to the specified number of times before failing the deploy. This is useful for applications that take a while to start up. The default is 7.
+You can also specify a custom healthcheck command, which is useful for non-HTTP services:
+
+```yaml
+healthcheck:
+  cmd: /bin/check_health
+```
+
+The top-level healthcheck configuration applies to all services that use
+Traefik, by default. You can also specialize the configuration at the role
+level:
+
+```yaml
+servers:
+  job:
+    hosts: ...
+    cmd: bin/jobs
+    healthcheck:
+      cmd: bin/check
+```
+
+The healthcheck allows for an optional `max_attempts` setting, which will attempt the healthcheck up to the specified number of times before failing the deploy. This is useful for applications that take a while to start up. The default is 7.
+
+Note that the HTTP health checks assume that the `curl` command is avilable inside the container. If that's not the case, use the healthcheck's `cmd` option to specify an alternative check that the container supports.
 
 ## Commands
 

lib/mrsk/cli/app.rb

@@ -6,8 +6,6 @@ class Mrsk::Cli::App < Mrsk::Cli::Base
       using_version(version_or_latest) do |version|
         say "Start container with version #{version} using a #{MRSK.config.readiness_delay}s readiness delay (or reboot if already running)...", :magenta
 
-        cli = self
-
         on(MRSK.hosts) do
           execute *MRSK.auditor.record("Tagging #{MRSK.config.absolute_image} as the latest image"), verbosity: :debug
           execute *MRSK.app.tag_current_as_latest
@@ -17,19 +15,24 @@ class Mrsk::Cli::App < Mrsk::Cli::Base
           roles = MRSK.roles_on(host)
 
           roles.each do |role|
-            execute *MRSK.auditor(role: role).record("Booted app version #{version}"), verbosity: :debug
+            app = MRSK.app(role: role)
+            auditor = MRSK.auditor(role: role)
 
-            if capture_with_info(*MRSK.app(role: role).container_id_for_version(version), raise_on_non_zero_exit: false).present?
+            execute *auditor.record("Booted app version #{version}"), verbosity: :debug
+
+            if capture_with_info(*app.container_id_for_version(version), raise_on_non_zero_exit: false).present?
               tmp_version = "#{version}_#{SecureRandom.hex(8)}"
               info "Renaming container #{version} to #{tmp_version} as already deployed on #{host}"
-              execute *MRSK.auditor(role: role).record("Renaming container #{version} to #{tmp_version}"), verbosity: :debug
+              execute *auditor.record("Renaming container #{version} to #{tmp_version}"), verbosity: :debug
-              execute *MRSK.app(role: role).rename_container(version: version, new_version: tmp_version)
+              execute *app.rename_container(version: version, new_version: tmp_version)
             end
 
-            old_version = capture_with_info(*MRSK.app(role: role).current_running_version, raise_on_non_zero_exit: false).strip
+            old_version = capture_with_info(*app.current_running_version, raise_on_non_zero_exit: false).strip
-            execute *MRSK.app(role: role).run
+            execute *app.run
-            sleep MRSK.config.readiness_delay
+
-            execute *MRSK.app(role: role).stop(version: old_version), raise_on_non_zero_exit: false if old_version.present?
+            Mrsk::Utils::HealthcheckPoller.wait_for_healthy(pause_after_ready: true) { capture_with_info(*app.status(version: version)) }
+
+            execute *app.stop(version: old_version), raise_on_non_zero_exit: false if old_version.present?
           end
         end
       end

lib/mrsk/cli/base.rb

@@ -107,7 +107,7 @@ module Mrsk::Cli
         MRSK.holding_lock = true
       rescue SSHKit::Runner::ExecuteError => e
         if e.message =~ /cannot create directory/
-          invoke "mrsk:cli:lock:status", []
+          on(MRSK.primary_host) { execute *MRSK.lock.status }
           raise LockError, "Deploy lock found"
         else
           raise e

lib/mrsk/cli/healthcheck.rb

@@ -1,7 +1,4 @@
 class Mrsk::Cli::Healthcheck < Mrsk::Cli::Base
-
-  class HealthcheckError < StandardError; end
-
   default_command :perform
 
   desc "perform", "Health check current app version"
@@ -9,38 +6,10 @@ class Mrsk::Cli::Healthcheck < Mrsk::Cli::Base
     on(MRSK.primary_host) do
       begin
         execute *MRSK.healthcheck.run
-
+        Mrsk::Utils::HealthcheckPoller.wait_for_healthy { capture_with_info(*MRSK.healthcheck.status) }
-        target = "Health check against #{MRSK.config.healthcheck["path"]}"
+      rescue Mrsk::Utils::HealthcheckPoller::HealthcheckError => e
-        attempt = 1
-        max_attempts = MRSK.config.healthcheck["max_attempts"]
-
-        begin
-          status = capture_with_info(*MRSK.healthcheck.curl)
-
-          if status == "200"
-            info "#{target} succeeded with 200 OK!"
-          else
-            raise HealthcheckError, "#{target} failed with status #{status}"
-          end
-        rescue SSHKit::Command::Failed
-          if attempt <= max_attempts
-            info "#{target} failed to respond, retrying in #{attempt}s (attempt #{attempt}/#{max_attempts})..."
-            sleep attempt
-            attempt += 1
-
-            retry
-          else
-            raise
-          end
-        end
-      rescue SSHKit::Command::Failed, HealthcheckError => e
         error capture_with_info(*MRSK.healthcheck.logs)
-
+        raise
-        if e.message =~ /curl/
-          raise SSHKit::Command::Failed, "#{target} failed to return 200 OK!"
-        else
-          raise
-        end
       ensure
         execute *MRSK.healthcheck.stop, raise_on_non_zero_exit: false
         execute *MRSK.healthcheck.remove, raise_on_non_zero_exit: false

lib/mrsk/cli/traefik.rb

@@ -94,7 +94,7 @@ class Mrsk::Cli::Traefik < Mrsk::Cli::Base
     end
   end
 
-  desc "remove_container", "Remove Traefik image from servers", hide: true
+  desc "remove_image", "Remove Traefik image from servers", hide: true
   def remove_image
     with_lock do
       on(MRSK.traefik_hosts) do

lib/mrsk/commands/app.rb

@@ -15,6 +15,7 @@ class Mrsk::Commands::App < Mrsk::Commands::Base
       "--name", container_name,
       "-e", "MRSK_CONTAINER_NAME=\"#{container_name}\"",
       *role.env_args,
+      *role.health_check_args,
       *config.logging_args,
       *config.volume_args,
       *role.label_args,
@@ -27,6 +28,10 @@ class Mrsk::Commands::App < Mrsk::Commands::Base
     docker :start, container_name
   end
 
+  def status(version:)
+    pipe container_id_for_version(version), xargs(docker(:inspect, "--format", DOCKER_HEALTH_STATUS_FORMAT))
+  end
+
   def stop(version: nil)
     pipe \
       version ? container_id_for_version(version) : current_running_container_id,

lib/mrsk/commands/base.rb

@@ -2,6 +2,8 @@ module Mrsk::Commands
   class Base
     delegate :sensitive, :argumentize, to: Mrsk::Utils
 
+    DOCKER_HEALTH_STATUS_FORMAT = "'{{if .State.Health}}{{.State.Health.Status}}{{else}}{{.State.Status}}{{end}}'"
+
     attr_accessor :config
 
     def initialize(config)

lib/mrsk/commands/healthcheck.rb

@@ -11,14 +11,15 @@ class Mrsk::Commands::Healthcheck < Mrsk::Commands::Base
       "--label", "service=#{container_name}",
       "-e", "MRSK_CONTAINER_NAME=\"#{container_name}\"",
       *web.env_args,
+      *web.health_check_args,
       *config.volume_args,
       *web.option_args,
       config.absolute_image,
       web.cmd
   end
 
-  def curl
+  def status
-    [ :curl, "--silent", "--output", "/dev/null", "--write-out", "'%{http_code}'", "--max-time", "2", health_url ]
+    pipe container_id, xargs(docker(:inspect, "--format", DOCKER_HEALTH_STATUS_FORMAT))
   end
 
   def logs

lib/mrsk/configuration/role.rb

@@ -35,6 +35,21 @@ class Mrsk::Configuration::Role
     argumentize_env_with_secrets env
   end
 
+  def health_check_args
+    if health_check_cmd.present?
+      optionize({ "health-cmd" => health_check_cmd, "health-interval" => "1s" })
+    else
+      []
+    end
+  end
+
+  def health_check_cmd
+    options = specializations["healthcheck"] || {}
+    options = config.healthcheck.merge(options) if running_traefik?
+
+    options["cmd"] || http_health_check(port: options["port"], path: options["path"])
+  end
+
   def cmd
     specializations["cmd"]
   end
@@ -75,8 +90,6 @@ class Mrsk::Configuration::Role
       if running_traefik?
         {
           "traefik.http.routers.#{traefik_service}.rule" => "PathPrefix(`/`)",
-          "traefik.http.services.#{traefik_service}.loadbalancer.healthcheck.path" => config.healthcheck["path"],
-          "traefik.http.services.#{traefik_service}.loadbalancer.healthcheck.interval" => "1s",
           "traefik.http.middlewares.#{traefik_service}-retry.retry.attempts" => "5",
           "traefik.http.middlewares.#{traefik_service}-retry.retry.initialinterval" => "500ms",
           "traefik.http.routers.#{traefik_service}.middlewares" => "#{traefik_service}-retry@docker"
@@ -125,4 +138,8 @@ class Mrsk::Configuration::Role
         new_env["clear"] = (clear_app_env + clear_role_env).uniq
       end
     end
+
+    def http_health_check(port:, path:)
+      "curl -f #{URI.join("http://localhost:#{port}", path)} || exit 1" if path.present? || port.present?
+    end
 end

lib/mrsk/utils/healthcheck_poller.rb

@@ -0,0 +1,39 @@
+class Mrsk::Utils::HealthcheckPoller
+  TRAEFIK_HEALTHY_DELAY = 1
+
+  class HealthcheckError < StandardError; end
+
+  class << self
+    def wait_for_healthy(pause_after_ready: false, &block)
+      attempt = 1
+      max_attempts = MRSK.config.healthcheck["max_attempts"]
+
+      begin
+        case status = block.call
+        when "healthy"
+          sleep TRAEFIK_HEALTHY_DELAY if pause_after_ready
+        when "running" # No health check configured
+          sleep MRSK.config.readiness_delay if pause_after_ready
+        else
+          raise HealthcheckError, "container not ready (#{status})"
+        end
+      rescue HealthcheckError => e
+        if attempt <= max_attempts
+          info "#{e.message}, retrying in #{attempt}s (attempt #{attempt}/#{max_attempts})..."
+          sleep attempt
+          attempt += 1
+          retry
+        else
+          raise
+        end
+      end
+
+      info "Container is healthy!"
+    end
+
+    private
+      def info(message)
+        SSHKit.config.output.info(message)
+      end
+  end
+end

test/cli/app_test.rb

@@ -2,8 +2,11 @@ require_relative "cli_test_case"
 
 class CliAppTest < CliTestCase
   test "boot" do
-    # Stub current version fetch
+    SSHKit::Backend::Abstract.any_instance.stubs(:capture_with_info).returns("123") # old version
-    SSHKit::Backend::Abstract.any_instance.stubs(:capture).returns("123") # old version
+
+    SSHKit::Backend::Abstract.any_instance.expects(:capture_with_info)
+    .with(:docker, :container, :ls, "--all", "--filter", "name=^app-web-latest$", "--quiet", "|", :xargs, :docker, :inspect, "--format", "'{{if .State.Health}}{{.State.Health.Status}}{{else}}{{.State.Status}}{{end}}'")
+    .returns("running") # health check
 
     run_command("boot").tap do |output|
       assert_match "docker tag dhh/app:latest dhh/app:latest", output
@@ -19,6 +22,10 @@ class CliAppTest < CliTestCase
       .with(:docker, :container, :ls, "--all", "--filter", "name=^app-web-latest$", "--quiet", raise_on_non_zero_exit: false)
       .returns("12345678") # running version
 
+    SSHKit::Backend::Abstract.any_instance.expects(:capture_with_info)
+      .with(:docker, :container, :ls, "--all", "--filter", "name=^app-web-latest$", "--quiet", "|", :xargs, :docker, :inspect, "--format", "'{{if .State.Health}}{{.State.Health.Status}}{{else}}{{.State.Status}}{{end}}'")
+      .returns("running") # health check
+
     SSHKit::Backend::Abstract.any_instance.expects(:capture_with_info)
       .with(:docker, :ps, "--filter", "label=service=app", "--filter", "label=role=web", "--filter", "status=running", "--latest", "--format", "\"{{.Names}}\"", "|", "grep -oE \"\\-[^-]+$\"", "|", "cut -c 2-", raise_on_non_zero_exit: false)
       .returns("123") # old version

test/cli/cli_test_case.rb

@@ -1,5 +1,4 @@
 require "test_helper"
-require "active_support/testing/stream"
 
 class CliTestCase < ActiveSupport::TestCase
   include ActiveSupport::Testing::Stream
@@ -17,13 +16,4 @@ class CliTestCase < ActiveSupport::TestCase
     ENV.delete("MYSQL_ROOT_PASSWORD")
     ENV.delete("VERSION")
   end
-
+end
-  private
-    def stdouted
-      capture(:stdout) { yield }.strip
-    end
-
-    def stderred
-      capture(:stderr) { yield }.strip
-    end
-  end

test/cli/healthcheck_test.rb

@@ -5,62 +5,58 @@ class CliHealthcheckTest < CliTestCase
     # Prevent expected failures from outputting to terminal
     Thread.report_on_exception = false
 
-    SSHKit::Backend::Abstract.any_instance.stubs(:sleep) # No sleeping when retrying
+    Mrsk::Utils::HealthcheckPoller.stubs(:sleep) # No sleeping when retrying
+
     SSHKit::Backend::Abstract.any_instance.stubs(:execute)
       .with(:docker, :container, :ls, "--all", "--filter", "name=^healthcheck-app-999$", "--quiet", "|", :xargs, :docker, :stop, raise_on_non_zero_exit: false)
     SSHKit::Backend::Abstract.any_instance.stubs(:execute)
-      .with(:docker, :run, "--detach", "--name", "healthcheck-app-999", "--publish", "3999:3000", "--label", "service=healthcheck-app", "-e", "MRSK_CONTAINER_NAME=\"healthcheck-app\"", "dhh/app:999")
+      .with(:docker, :run, "--detach", "--name", "healthcheck-app-999", "--publish", "3999:3000", "--label", "service=healthcheck-app", "-e", "MRSK_CONTAINER_NAME=\"healthcheck-app\"", "--health-cmd", "\"curl -f http://localhost:3000/up || exit 1\"", "--health-interval", "\"1s\"", "dhh/app:999")
     SSHKit::Backend::Abstract.any_instance.stubs(:execute)
       .with(:docker, :container, :ls, "--all", "--filter", "name=^healthcheck-app-999$", "--quiet", "|", :xargs, :docker, :container, :rm, raise_on_non_zero_exit: false)
 
     # Fail twice to test retry logic
     SSHKit::Backend::Abstract.any_instance.stubs(:capture_with_info)
-      .with(:curl, "--silent", "--output", "/dev/null", "--write-out", "'%{http_code}'", "--max-time", "2", "http://localhost:3999/up")
+      .with(:docker, :container, :ls, "--all", "--filter", "name=^healthcheck-app-999$", "--quiet", "|", :xargs, :docker, :inspect, "--format", "'{{if .State.Health}}{{.State.Health.Status}}{{else}}{{.State.Status}}{{end}}'")
-      .raises(SSHKit::Command::Failed)
+      .returns("starting")
       .then
-      .raises(SSHKit::Command::Failed)
+      .returns("unhealthy")
       .then
-      .returns("200")
+      .returns("healthy")
 
     run_command("perform").tap do |output|
-      assert_match "Health check against /up failed to respond, retrying in 1s (attempt 1/7)...", output
+      assert_match "container not ready (starting), retrying in 1s (attempt 1/7)...", output
-      assert_match "Health check against /up failed to respond, retrying in 2s (attempt 2/7)...", output
+      assert_match "container not ready (unhealthy), retrying in 2s (attempt 2/7)...", output
-      assert_match "Health check against /up succeeded with 200 OK!", output
+      assert_match "Container is healthy!", output
     end
   end
 
-  test "perform failing because of curl" do
+  test "perform failing to become healthy" do
     # Prevent expected failures from outputting to terminal
     Thread.report_on_exception = false
 
-    SSHKit::Backend::Abstract.any_instance.stubs(:execute) # No need to execute anything here
+    Mrsk::Utils::HealthcheckPoller.stubs(:sleep) # No sleeping when retrying
+
+    SSHKit::Backend::Abstract.any_instance.stubs(:execute)
+      .with(:docker, :container, :ls, "--all", "--filter", "name=^healthcheck-app-999$", "--quiet", "|", :xargs, :docker, :stop, raise_on_non_zero_exit: false)
+    SSHKit::Backend::Abstract.any_instance.stubs(:execute)
+      .with(:docker, :run, "--detach", "--name", "healthcheck-app-999", "--publish", "3999:3000", "--label", "service=healthcheck-app", "-e", "MRSK_CONTAINER_NAME=\"healthcheck-app\"", "--health-cmd", "\"curl -f http://localhost:3000/up || exit 1\"", "--health-interval", "\"1s\"", "dhh/app:999")
+    SSHKit::Backend::Abstract.any_instance.stubs(:execute)
+      .with(:docker, :container, :ls, "--all", "--filter", "name=^healthcheck-app-999$", "--quiet", "|", :xargs, :docker, :container, :rm, raise_on_non_zero_exit: false)
+
+    # Continually report unhealthy
     SSHKit::Backend::Abstract.any_instance.stubs(:capture_with_info)
-      .with(:curl, "--silent", "--output", "/dev/null", "--write-out", "'%{http_code}'", "--max-time", "2", "http://localhost:3999/up")
+      .with(:docker, :container, :ls, "--all", "--filter", "name=^healthcheck-app-999$", "--quiet", "|", :xargs, :docker, :inspect, "--format", "'{{if .State.Health}}{{.State.Health.Status}}{{else}}{{.State.Status}}{{end}}'")
-      .returns("curl: command not found")
+      .returns("unhealthy")
-    SSHKit::Backend::Abstract.any_instance.stubs(:capture_with_info)
+
-      .with(:docker, :container, :ls, "--all", "--filter", "name=^healthcheck-app-999$", "--quiet", "|", :xargs, :docker, :logs, "--tail", 50, "2>&1")
+    # Capture logs when failing
-
-    exception = assert_raises SSHKit::Runner::ExecuteError do
-      run_command("perform")
-    end
-    assert_match "Health check against /up failed to return 200 OK!", exception.message
-  end
-
-  test "perform failing for unknown reason" do
-    # Prevent expected failures from outputting to terminal
-    Thread.report_on_exception = false
-
-    SSHKit::Backend::Abstract.any_instance.stubs(:execute) # No need to execute anything here
-    SSHKit::Backend::Abstract.any_instance.stubs(:capture_with_info)
-      .with(:curl, "--silent", "--output", "/dev/null", "--write-out", "'%{http_code}'", "--max-time", "2", "http://localhost:3999/up")
-      .returns("500")
     SSHKit::Backend::Abstract.any_instance.stubs(:capture_with_info)
       .with(:docker, :container, :ls, "--all", "--filter", "name=^healthcheck-app-999$", "--quiet", "|", :xargs, :docker, :logs, "--tail", 50, "2>&1")
+      .returns("some log output")
 
     exception = assert_raises do
       run_command("perform")
     end
-    assert_match "Health check against /up failed with status 500", exception.message
+    assert_match "container not ready (unhealthy)", exception.message
   end
 
   private

test/cli/main_test.rb

@@ -64,7 +64,8 @@ class CliMainTest < CliTestCase
       .with { |*arg| arg[0..1] == [:mkdir, :mrsk_lock] }
       .raises(RuntimeError, "mkdir: cannot create directory ‘mrsk_lock’: File exists")
 
-    Mrsk::Cli::Base.any_instance.expects(:invoke).with("mrsk:cli:lock:status", [])
+    SSHKit::Backend::Abstract.any_instance.expects(:execute)
+      .with(:stat, :mrsk_lock, ">", "/dev/null", "&&", :cat, "mrsk_lock/details", "|", :base64, "-d")
 
     assert_raises(Mrsk::Cli::LockError) do
       run_command("deploy")

test/commands/app_test.rb

@@ -13,7 +13,7 @@ class CommandsAppTest < ActiveSupport::TestCase
 
   test "run" do
     assert_equal \
-      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --log-opt max-size=\"10m\" --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.services.app-web.loadbalancer.healthcheck.path=\"/up\" --label traefik.http.services.app-web.loadbalancer.healthcheck.interval=\"1s\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
+      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --health-cmd \"curl -f http://localhost:3000/up || exit 1\" --health-interval \"1s\" --log-opt max-size=\"10m\" --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
       new_command.run.join(" ")
   end
 
@@ -21,7 +21,7 @@ class CommandsAppTest < ActiveSupport::TestCase
     @config[:volumes] = ["/local/path:/container/path" ]
 
     assert_equal \
-      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --log-opt max-size=\"10m\" --volume /local/path:/container/path --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.services.app-web.loadbalancer.healthcheck.path=\"/up\" --label traefik.http.services.app-web.loadbalancer.healthcheck.interval=\"1s\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
+      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --health-cmd \"curl -f http://localhost:3000/up || exit 1\" --health-interval \"1s\" --log-opt max-size=\"10m\" --volume /local/path:/container/path --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
       new_command.run.join(" ")
   end
 
@@ -29,7 +29,23 @@ class CommandsAppTest < ActiveSupport::TestCase
     @config[:healthcheck] = { "path" => "/healthz" }
 
     assert_equal \
-      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --log-opt max-size=\"10m\" --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.services.app-web.loadbalancer.healthcheck.path=\"/healthz\" --label traefik.http.services.app-web.loadbalancer.healthcheck.interval=\"1s\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
+      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --health-cmd \"curl -f http://localhost:3000/healthz || exit 1\" --health-interval \"1s\" --log-opt max-size=\"10m\" --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
+      new_command.run.join(" ")
+  end
+
+  test "run with custom healthcheck command" do
+    @config[:healthcheck] = { "cmd" => "/bin/up" }
+
+    assert_equal \
+      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --health-cmd \"/bin/up\" --health-interval \"1s\" --log-opt max-size=\"10m\" --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
+      new_command.run.join(" ")
+  end
+
+  test "run with role-specific healthcheck options" do
+    @config[:servers] = { "web" => { "hosts" => [ "1.1.1.1" ], "healthcheck" => { "cmd" => "/bin/healthy" } } }
+
+    assert_equal \
+      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --health-cmd \"/bin/healthy\" --health-interval \"1s\" --log-opt max-size=\"10m\" --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
       new_command.run.join(" ")
   end
 
@@ -44,7 +60,7 @@ class CommandsAppTest < ActiveSupport::TestCase
     @config[:logging] = { "driver" => "local", "options" => { "max-size" => "100m", "max-file" => "3" } }
 
     assert_equal \
-      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --log-driver \"local\" --log-opt max-size=\"100m\" --log-opt max-file=\"3\" --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.services.app-web.loadbalancer.healthcheck.path=\"/up\" --label traefik.http.services.app-web.loadbalancer.healthcheck.interval=\"1s\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
+      "docker run --detach --restart unless-stopped --name app-web-999 -e MRSK_CONTAINER_NAME=\"app-web-999\" -e RAILS_MASTER_KEY=\"456\" --health-cmd \"curl -f http://localhost:3000/up || exit 1\" --health-interval \"1s\" --log-driver \"local\" --log-opt max-size=\"100m\" --log-opt max-file=\"3\" --label service=\"app\" --label role=\"web\" --label traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\" --label traefik.http.middlewares.app-web-retry.retry.attempts=\"5\" --label traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\" --label traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\" dhh/app:999",
       new_command.run.join(" ")
   end
 

test/commands/healthcheck_test.rb

@@ -10,7 +10,7 @@ class CommandsHealthcheckTest < ActiveSupport::TestCase
 
   test "run" do
     assert_equal \
-      "docker run --detach --name healthcheck-app-123 --publish 3999:3000 --label service=healthcheck-app -e MRSK_CONTAINER_NAME=\"healthcheck-app\" dhh/app:123",
+      "docker run --detach --name healthcheck-app-123 --publish 3999:3000 --label service=healthcheck-app -e MRSK_CONTAINER_NAME=\"healthcheck-app\" --health-cmd \"curl -f http://localhost:3000/up || exit 1\" --health-interval \"1s\" dhh/app:123",
       new_command.run.join(" ")
   end
 
@@ -18,7 +18,7 @@ class CommandsHealthcheckTest < ActiveSupport::TestCase
     @config[:healthcheck] = { "port" => 3001 }
 
     assert_equal \
-      "docker run --detach --name healthcheck-app-123 --publish 3999:3001 --label service=healthcheck-app -e MRSK_CONTAINER_NAME=\"healthcheck-app\" dhh/app:123",
+      "docker run --detach --name healthcheck-app-123 --publish 3999:3001 --label service=healthcheck-app -e MRSK_CONTAINER_NAME=\"healthcheck-app\" --health-cmd \"curl -f http://localhost:3001/up || exit 1\" --health-interval \"1s\" dhh/app:123",
       new_command.run.join(" ")
   end
 
@@ -26,29 +26,29 @@ class CommandsHealthcheckTest < ActiveSupport::TestCase
     @destination = "staging"
 
     assert_equal \
-      "docker run --detach --name healthcheck-app-staging-123 --publish 3999:3000 --label service=healthcheck-app-staging -e MRSK_CONTAINER_NAME=\"healthcheck-app-staging\" dhh/app:123",
+      "docker run --detach --name healthcheck-app-staging-123 --publish 3999:3000 --label service=healthcheck-app-staging -e MRSK_CONTAINER_NAME=\"healthcheck-app-staging\" --health-cmd \"curl -f http://localhost:3000/up || exit 1\" --health-interval \"1s\" dhh/app:123",
+      new_command.run.join(" ")
+  end
+
+  test "run with custom healthcheck" do
+    @config[:healthcheck] = { "cmd" => "/bin/up" }
+
+    assert_equal \
+      "docker run --detach --name healthcheck-app-123 --publish 3999:3000 --label service=healthcheck-app -e MRSK_CONTAINER_NAME=\"healthcheck-app\" --health-cmd \"/bin/up\" --health-interval \"1s\" dhh/app:123",
       new_command.run.join(" ")
   end
 
   test "run with custom options" do
     @config[:servers] = { "web" => { "hosts" => [ "1.1.1.1" ], "options" => { "mount" => "somewhere" } } }
     assert_equal \
-      "docker run --detach --name healthcheck-app-123 --publish 3999:3000 --label service=healthcheck-app -e MRSK_CONTAINER_NAME=\"healthcheck-app\" --mount \"somewhere\" dhh/app:123",
+      "docker run --detach --name healthcheck-app-123 --publish 3999:3000 --label service=healthcheck-app -e MRSK_CONTAINER_NAME=\"healthcheck-app\" --health-cmd \"curl -f http://localhost:3000/up || exit 1\" --health-interval \"1s\" --mount \"somewhere\" dhh/app:123",
       new_command.run.join(" ")
   end
 
-  test "curl" do
+  test "status" do
     assert_equal \
-      "curl --silent --output /dev/null --write-out '%{http_code}' --max-time 2 http://localhost:3999/up",
+      "docker container ls --all --filter name=^healthcheck-app-123$ --quiet | xargs docker inspect --format '{{if .State.Health}}{{.State.Health.Status}}{{else}}{{.State.Status}}{{end}}'",
-      new_command.curl.join(" ")
+      new_command.status.join(" ")
-  end
-
-  test "curl with custom path" do
-    @config[:healthcheck] = { "path" => "/healthz" }
-
-    assert_equal \
-      "curl --silent --output /dev/null --write-out '%{http_code}' --max-time 2 http://localhost:3999/healthz",
-      new_command.curl.join(" ")
   end
 
   test "stop" do

test/configuration/role_test.rb

@@ -42,7 +42,7 @@ class ConfigurationRoleTest < ActiveSupport::TestCase
   end
 
   test "special label args for web" do
-    assert_equal [ "--label", "service=\"app\"", "--label", "role=\"web\"", "--label", "traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\"", "--label", "traefik.http.services.app-web.loadbalancer.healthcheck.path=\"/up\"", "--label", "traefik.http.services.app-web.loadbalancer.healthcheck.interval=\"1s\"", "--label", "traefik.http.middlewares.app-web-retry.retry.attempts=\"5\"", "--label", "traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\"", "--label", "traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\"" ], @config.role(:web).label_args
+    assert_equal [ "--label", "service=\"app\"", "--label", "role=\"web\"", "--label", "traefik.http.routers.app-web.rule=\"PathPrefix(\\`/\\`)\"", "--label", "traefik.http.middlewares.app-web-retry.retry.attempts=\"5\"", "--label", "traefik.http.middlewares.app-web-retry.retry.initialinterval=\"500ms\"", "--label", "traefik.http.routers.app-web.middlewares=\"app-web-retry@docker\"" ], @config.role(:web).label_args
   end
 
   test "custom labels" do
@@ -57,8 +57,8 @@ class ConfigurationRoleTest < ActiveSupport::TestCase
   end
 
   test "overwriting default traefik label" do
-    @deploy[:labels] = { "traefik.http.routers.app.rule" => "\"Host(\\`example.com\\`) || (Host(\\`example.org\\`) && Path(\\`/traefik\\`))\"" }
+    @deploy[:labels] = { "traefik.http.routers.app-web.rule" => "\"Host(\\`example.com\\`) || (Host(\\`example.org\\`) && Path(\\`/traefik\\`))\"" }
-    assert_equal "\"Host(\\`example.com\\`) || (Host(\\`example.org\\`) && Path(\\`/traefik\\`))\"", @config.role(:web).labels["traefik.http.routers.app.rule"]
+    assert_equal "\"Host(\\`example.com\\`) || (Host(\\`example.org\\`) && Path(\\`/traefik\\`))\"", @config.role(:web).labels["traefik.http.routers.app-web.rule"]
   end
 
   test "default traefik label on non-web role" do
@@ -66,15 +66,7 @@ class ConfigurationRoleTest < ActiveSupport::TestCase
       c[:servers]["beta"] = { "traefik" => "true", "hosts" => [ "1.1.1.5" ] }
     })
 
-    assert_equal [ "--label", "service=\"app\"", "--label", "role=\"beta\"", "--label", "traefik.http.routers.app-beta.rule=\"PathPrefix(\\`/\\`)\"", "--label", "traefik.http.services.app-beta.loadbalancer.healthcheck.path=\"/up\"", "--label", "traefik.http.services.app-beta.loadbalancer.healthcheck.interval=\"1s\"", "--label", "traefik.http.middlewares.app-beta-retry.retry.attempts=\"5\"", "--label", "traefik.http.middlewares.app-beta-retry.retry.initialinterval=\"500ms\"", "--label", "traefik.http.routers.app-beta.middlewares=\"app-beta-retry@docker\"" ], config.role(:beta).label_args
+    assert_equal [ "--label", "service=\"app\"", "--label", "role=\"beta\"", "--label", "traefik.http.routers.app-beta.rule=\"PathPrefix(\\`/\\`)\"", "--label", "traefik.http.middlewares.app-beta-retry.retry.attempts=\"5\"", "--label", "traefik.http.middlewares.app-beta-retry.retry.initialinterval=\"500ms\"", "--label", "traefik.http.routers.app-beta.middlewares=\"app-beta-retry@docker\"" ], config.role(:beta).label_args
-  end
-
-  test "default traefik label for non-web role with destination" do
-    config = Mrsk::Configuration.new(@deploy_with_roles.tap { |c|
-      c[:servers]["beta"] = { "traefik" => "true", "hosts" => [ "1.1.1.5" ] }
-    }, destination: "staging")
-
-    assert_equal [ "--label", "service=\"app\"", "--label", "role=\"beta\"", "--label", "destination=\"staging\"", "--label", "traefik.http.routers.app-beta-staging.rule=\"PathPrefix(\\`/\\`)\"", "--label", "traefik.http.services.app-beta-staging.loadbalancer.healthcheck.path=\"/up\"", "--label", "traefik.http.services.app-beta-staging.loadbalancer.healthcheck.interval=\"1s\"", "--label", "traefik.http.middlewares.app-beta-staging-retry.retry.attempts=\"5\"", "--label", "traefik.http.middlewares.app-beta-staging-retry.retry.initialinterval=\"500ms\"", "--label", "traefik.http.routers.app-beta-staging.middlewares=\"app-beta-staging-retry@docker\"" ], config.role(:beta).label_args
   end
 
   test "env overwritten by role" do

test/integration/deploy_test.rb

@@ -0,0 +1,68 @@
+require "net/http"
+require "test_helper"
+
+class DeployTest < ActiveSupport::TestCase
+
+  setup do
+    docker_compose "up --build --force-recreate -d"
+    wait_for_healthy
+  end
+
+  teardown do
+    docker_compose "down -v"
+  end
+
+  test "deploy" do
+    assert_app_is_down
+
+    mrsk :deploy
+
+    assert_app_is_up
+  end
+
+  private
+    def docker_compose(*commands, capture: false)
+      command = "docker compose #{commands.join(" ")}"
+      succeeded = false
+      if capture
+        result = stdouted { succeeded = system("cd test/integration && #{command}") }
+      else
+        succeeded = system("cd test/integration && #{command}")
+      end
+
+      raise "Command `#{command}` failed with error code `#{$?}`" unless succeeded
+      result
+    end
+
+    def deployer_exec(*commands, capture: false)
+      if capture
+        stdouted { docker_compose("exec deployer #{commands.join(" ")}") }
+      else
+        docker_compose("exec deployer #{commands.join(" ")}", capture: capture)
+      end
+    end
+
+    def mrsk(*commands, capture: false)
+      deployer_exec(:mrsk, *commands, capture: capture)
+    end
+
+    def assert_app_is_down
+      assert_equal "502", app_response.code
+    end
+
+    def assert_app_is_up
+      assert_equal "200", app_response.code
+    end
+
+    def app_response
+      Net::HTTP.get_response(URI.parse("http://localhost:12345"))
+    end
+
+    def wait_for_healthy(timeout: 20)
+      timeout_at = Time.now + timeout
+      while docker_compose("ps -a | tail -n +2 | grep -v '(healthy)' | wc -l", capture: true) != "0"
+        raise "Container not healthy after #{timeout} seconds" if timeout_at < Time.now
+        sleep 0.1
+      end
+    end
+end

test/integration/docker-compose.yml

@@ -0,0 +1,50 @@
+version: "3.7"
+name: "mrsk-test"
+
+volumes:
+  shared:
+
+services:
+  shared:
+    build:
+      context: docker/shared
+    volumes:
+      - shared:/shared
+
+  deployer:
+    privileged: true
+    build:
+      context: docker/deployer
+    volumes:
+      - ../..:/mrsk
+      - shared:/shared
+
+  registry:
+    build:
+      context: docker/registry
+    environment:
+      - REGISTRY_HTTP_ADDR=0.0.0.0:4443
+      - REGISTRY_HTTP_TLS_CERTIFICATE=/certs/domain.crt
+      - REGISTRY_HTTP_TLS_KEY=/certs/domain.key
+    volumes:
+      - shared:/shared
+
+  vm1:
+    privileged: true
+    build:
+      context: docker/vm
+    volumes:
+      - shared:/shared
+
+  vm2:
+    privileged: true
+    build:
+      context: docker/vm
+    volumes:
+      - shared:/shared
+
+  load_balancer:
+    build:
+      context: docker/load_balancer
+    ports:
+      - "12345:80"

test/integration/docker/deployer/Dockerfile

@@ -0,0 +1,29 @@
+FROM ruby:3.2
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y ca-certificates openssh-client curl gnupg docker.io
+
+RUN install -m 0755 -d /etc/apt/keyrings
+RUN curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg
+RUN chmod a+r /etc/apt/keyrings/docker.gpg
+RUN echo \
+  "deb [arch="$(dpkg --print-architecture)" signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/debian \
+  "$(. /etc/os-release && echo "$VERSION_CODENAME")" stable" | \
+  tee /etc/apt/sources.list.d/docker.list > /dev/null
+
+RUN apt-get update && apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
+
+COPY boot.sh .
+COPY app/ .
+
+RUN ln -s /shared/ssh /root/.ssh
+RUN mkdir -p /etc/docker/certs.d/registry:4443 && ln -s /shared/certs/domain.crt /etc/docker/certs.d/registry:4443/ca.crt
+
+RUN git config --global user.email "deployer@example.com"
+RUN git config --global user.name "Deployer"
+RUN git init && git add . && git commit -am "Initial version"
+
+HEALTHCHECK --interval=1s CMD pgrep sleep
+
+CMD ["./boot.sh"]

test/integration/docker/deployer/app/Dockerfile

@@ -0,0 +1,3 @@
+FROM nginx:1-alpine-slim
+
+COPY default.conf /etc/nginx/conf.d/default.conf

test/integration/docker/deployer/app/config/deploy.yml

@@ -0,0 +1,14 @@
+service: app
+image: app
+servers:
+  - vm1
+  - vm2
+registry:
+  server: registry:4443
+  username: root
+  password: root
+builder:
+  multiarch: false
+healthcheck:
+  path: /
+  port: 80

test/integration/docker/deployer/app/default.conf

@@ -0,0 +1,17 @@
+server {
+    listen       80;
+    listen  [::]:80;
+    server_name  localhost;
+
+    location / {
+        root   /usr/share/nginx/html;
+        index  index.html index.htm;
+    }
+
+    # redirect server error pages to the static page /50x.html
+    #
+    error_page   500 502 503 504  /50x.html;
+    location = /50x.html {
+        root   /usr/share/nginx/html;
+    }
+}

test/integration/docker/deployer/boot.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+cd /mrsk && gem build mrsk.gemspec -o /tmp/mrsk.gem && gem install /tmp/mrsk.gem
+
+dockerd &
+
+trap "pkill -f sleep" term
+
+sleep infinity & wait

test/integration/docker/load_balancer/Dockerfile

@@ -0,0 +1,5 @@
+FROM nginx:1-alpine-slim
+
+COPY default.conf /etc/nginx/conf.d/default.conf
+
+HEALTHCHECK --interval=1s CMD pgrep nginx

test/integration/docker/load_balancer/default.conf

@@ -0,0 +1,12 @@
+upstream loadbalancer {
+  server vm1:80;
+  server vm2:80;
+}
+
+server {
+  listen 80;
+
+  location / {
+    proxy_pass http://loadbalancer;
+  }
+}

test/integration/docker/registry/Dockerfile

@@ -0,0 +1,9 @@
+FROM registry
+
+COPY boot.sh .
+
+RUN ln -s /shared/certs /certs
+
+HEALTHCHECK --interval=1s CMD pgrep registry
+
+ENTRYPOINT ["./boot.sh"]

test/integration/docker/registry/boot.sh

@@ -0,0 +1,7 @@
+#!/bin/sh
+
+while [ ! -f /certs/domain.crt ]; do sleep 1; done
+
+trap "pkill -f registry" term
+
+/entrypoint.sh /etc/docker/registry/config.yml & wait

test/integration/docker/shared/Dockerfile

@@ -0,0 +1,17 @@
+FROM ubuntu:22.10
+
+WORKDIR /work
+
+RUN apt-get update && apt-get -y install openssh-client openssl
+
+RUN mkdir ssh && \
+  ssh-keygen -t rsa -f ssh/id_rsa -N ""
+
+COPY registry-dns.conf .
+COPY boot.sh .
+
+RUN mkdir certs && openssl req -newkey rsa:4096 -nodes -sha256 -keyout certs/domain.key   -x509 -days 365 -out certs/domain.crt   -subj '/CN=registry' -extensions EXT -config registry-dns.conf
+
+HEALTHCHECK --interval=1s CMD pgrep sleep
+
+CMD ["./boot.sh"]

test/integration/docker/shared/boot.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+cp -r * /shared
+
+trap "pkill -f sleep" term
+
+sleep infinity & wait

test/integration/docker/shared/registry-dns.conf

@@ -0,0 +1,7 @@
+[dn]
+CN=registry
+[req]
+distinguished_name = dn
+[EXT]
+subjectAltName=DNS:registry
+keyUsage=digitalSignature

test/integration/docker/vm/Dockerfile

@@ -0,0 +1,14 @@
+FROM ubuntu:22.10
+
+WORKDIR /work
+
+RUN apt-get update && apt-get -y install openssh-client openssh-server docker.io
+
+RUN mkdir /root/.ssh && ln -s /shared/ssh/id_rsa.pub /root/.ssh/authorized_keys
+RUN mkdir -p /etc/docker/certs.d/registry:4443 && ln -s /shared/certs/domain.crt /etc/docker/certs.d/registry:4443/ca.crt
+
+COPY boot.sh .
+
+HEALTHCHECK --interval=1s CMD pgrep dockerd
+
+CMD ["./boot.sh"]

test/integration/docker/vm/boot.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+while [ ! -f /root/.ssh/authorized_keys ]; do echo "Waiting for ssh keys"; sleep 1; done
+
+service ssh restart
+
+dockerd &
+
+trap "pkill -f sleep" term
+
+sleep infinity & wait

test/test_helper.rb

@@ -1,6 +1,7 @@
 require "bundler/setup"
 require "active_support/test_case"
 require "active_support/testing/autorun"
+require "active_support/testing/stream"
 require "debug"
 require "mocha/minitest" # using #stubs that can alter returns
 require "minitest/autorun" # using #stub that take args
@@ -23,4 +24,14 @@ module SSHKit
 end
 
 class ActiveSupport::TestCase
+  include ActiveSupport::Testing::Stream
+
+  private
+    def stdouted
+      capture(:stdout) { yield }.strip
+    end
+
+    def stderred
+      capture(:stderr) { yield }.strip
+    end
 end