From ab6b960861c09a3aac546742b678a63c6f2d842d Mon Sep 17 00:00:00 2001
From: Cory Johns <johnsca@gmail.com>
Date: Tue, 14 Jun 2016 15:31:08 -0400
Subject: [PATCH] Fixes for dhx to work with 2.0

---
 juju-dhx | 42 ++++++++++++++++++++++++++++++------------
 1 file changed, 30 insertions(+), 12 deletions(-)

diff --git a/juju-dhx b/juju-dhx
index 2406706..33d6810 100755
--- a/juju-dhx
+++ b/juju-dhx
@@ -24,6 +24,19 @@ DEFAULT_CONFIG = {'import_ids': [],
                   'sync_exclude': ['*.pyc', '.*']}
 
 
+def call(*args):
+    try:
+        return subprocess.check_output(('/usr/bin/env',) + args,
+                                       stderr=subprocess.STDOUT)
+    except subprocess.CalledProcessError as e:
+        sys.stderr.write(e.output)
+        sys.exit(e.returncode)
+
+
+JUJU_VERSION = call('juju', 'version')
+JUJU_2 = JUJU_VERSION.startswith('2.')
+
+
 def show_help(option, opt, value, parser):
     print 'usage: juju dhx [options] [<unit-or-service>]'
     print 'purpose: Enhanced interactive remote hook debugging session, using tmux'
@@ -172,7 +185,13 @@ def do_import_ids(unit, import_ids, opts):
     machine = unit.get('Machine')
     annotations = env.get_annotation(machine, 'machine')
     if annotations['Annotations'].get('import-ids') != ','.join(import_ids):
-        call('juju', 'authorized-keys', 'import', *import_ids)
+        cmd = ['juju']
+        if JUJU_2:
+            cmd.append('import-ssh-keys')
+        else:
+            cmd.extend(['authorized-keys', 'import'])
+        cmd.extend(import_ids)
+        call(*cmd)
         env.set_annotation(machine, 'machine', {'import-ids': ','.join(import_ids)})
 
 
@@ -183,20 +202,18 @@ def get_env(opts):
     return get_env._env
 
 
-def call(*args):
-    try:
-        return subprocess.check_output(('/usr/bin/env',) + args,
-                                       stderr=subprocess.STDOUT)
-    except subprocess.CalledProcessError as e:
-        sys.stderr.write(e.output)
-        sys.exit(e.returncode)
+def _unit_errored(unit):
+    if JUJU_2:
+        return unit['WorkloadStatus']['Status'] == 'error'
+    else:
+        return unit['AgentState'] == 'error'
 
 
 def find_errored_unit(units):
     if not units:
         return None
     for name, unit in units.items():
-        if unit['AgentState'] == 'error':
+        if _unit_errored(unit):
             return name
     return None
 
@@ -259,7 +276,7 @@ def choose_unit_from(all_units, preferred_units=None):
     print 'Units:'
     for i, unit_name in enumerate(unit_names):
         unit_err = ''
-        if all_units[unit_name]['AgentState'] == 'error':
+        if _unit_errored(all_units[unit_name]):
             unit_err = ' (error)'
         print '  %d: %s%s' % (i, unit_name, unit_err)
     print
@@ -327,12 +344,13 @@ if __name__ == '__main__':
     else:
         unit_name, unit = choose_unit(args, opts)
         config = customize(unit_name, unit, opts)
+        model_opt = '-m' if JUJU_2 else '-e'
         if (opts.restart or config.get('auto_restart', False)) and not opts.no_restart:
             subprocess.Popen(
-                'sleep 2 ; juju resolved --retry -e %s "%s"' % (opts.env, unit_name),
+                'sleep 2 ; juju resolved --retry %s %s "%s"' % (model_opt, opts.env, unit_name),
                 stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True)
         result = subprocess.call(
-            ['/usr/bin/env', 'juju', 'debug-hooks', '-e', opts.env, unit_name] + args)
+            ['/usr/bin/env', 'juju', 'debug-hooks', model_opt, opts.env, unit_name] + args)
         if result != 0:
             sys.exit(result)
         if (opts.sync or config.get('auto_sync', False)) and not opts.no_sync: