From 1d662dc8f9ce487d332c0bf0b899095da03d7d9d Mon Sep 17 00:00:00 2001
From: Matt Brictson <matt@mattbrictson.com>
Date: Fri, 30 Oct 2015 15:22:32 -0700
Subject: [PATCH] Fall back to ssh_options[:user] if host.user.nil?

Up until now Airbrussh has displayed the username using `host.name`. However,
this may be `nil` in favor of using `ssh_options[:user]`. Use the SSH options
in this case.

If for some reason no user is specified with either method, then omit the `@`
symbol and just print the hostname instead of `@hostname`.

Addresses #65.
---
 CHANGELOG.md                             |  3 +++
 lib/airbrussh/command_formatter.rb       |  6 ++---
 test/airbrussh/command_formatter_test.rb | 34 ++++++++++++++++++++----
 3 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 09e75d6..e8d07a5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,9 @@ Airbrussh is in a pre-1.0 state. This means that its APIs and behavior are subje
 ## [Unreleased]
 
 * Your contribution here!
+* Airbrussh now displays the correct user@host output in the following edge-cases:
+    * Inside an SSHKit `as(:user => "...")` block
+    * When a user is specified using `:ssh_options => { :user => "..." }` ([see #65](https://github.com/mattbrictson/airbrussh/issues/65))
 
 ## [0.7.0][] (2015-08-08)
 
diff --git a/lib/airbrussh/command_formatter.rb b/lib/airbrussh/command_formatter.rb
index 2652764..e83a336 100644
--- a/lib/airbrussh/command_formatter.rb
+++ b/lib/airbrussh/command_formatter.rb
@@ -58,9 +58,9 @@ def exit_message(log_file=nil)
     private
 
     def user_at_host
-      user_str = user { host.user }
-      host_str = host.to_s
-      [user_str, host_str].join("@")
+      user_str = host.user || (host.ssh_options || {})[:user]
+      host_str = host.hostname
+      [user_str, host_str].compact.join("@")
     end
 
     def runtime
diff --git a/test/airbrussh/command_formatter_test.rb b/test/airbrussh/command_formatter_test.rb
index 3429ddb..ce60d8d 100644
--- a/test/airbrussh/command_formatter_test.rb
+++ b/test/airbrussh/command_formatter_test.rb
@@ -5,16 +5,16 @@
 
 class Airbrussh::CommandFormatterTest < Minitest::Test
   def setup
-    sshkit_command = OpenStruct.new(
-      :host => "12.34.56.78",
-      :user => "deployer",
+    @sshkit_command = OpenStruct.new(
+      :host => host("deployer", "12.34.56.78"),
+      :options => { :user => "override" },
       :runtime => 1.23456,
       :failure? => false
     )
-    def sshkit_command.to_s
+    def @sshkit_command.to_s
       "/usr/bin/env echo hello"
     end
-    @command = Airbrussh::CommandFormatter.new(sshkit_command, 0)
+    @command = Airbrussh::CommandFormatter.new(@sshkit_command, 0)
   end
 
   def test_format_output
@@ -39,4 +39,28 @@ def test_exit_message_failure
         @command.exit_message("out.log"))
     end
   end
+
+  def test_uses_ssh_options_if_host_user_is_absent
+    @sshkit_command.host = host(nil, "12.34.56.78", :user => "sshuser")
+    assert_equal(
+      "\e[0;32;49m✔ 01 sshuser@12.34.56.78\e[0m \e[0;90;49m1.235s\e[0m",
+      @command.exit_message)
+  end
+
+  def test_shows_hostname_only_if_no_user
+    @sshkit_command.host = host(nil, "12.34.56.78")
+    assert_equal(
+      "\e[0;32;49m✔ 01 12.34.56.78\e[0m \e[0;90;49m1.235s\e[0m",
+      @command.exit_message)
+  end
+
+  private
+
+  def host(user, hostname, ssh_options={})
+    SSHKit::Host.new(
+      :user => user,
+      :hostname => hostname,
+      :ssh_options => ssh_options
+    )
+  end
 end