From fb837cc2a630bb407cf059dc1871a2cc832cf801 Mon Sep 17 00:00:00 2001
From: Caner Derici <caner.derici@canonical.com>
Date: Wed, 10 Aug 2022 17:18:15 -0600
Subject: [PATCH] Fix for unit.run_actions to support both old and new clients

Fixes #705
---
 juju/action.py |  8 +++++++-
 juju/model.py  | 23 +++++++++++++++--------
 juju/unit.py   | 15 ++++++++++-----
 3 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/juju/action.py b/juju/action.py
index 068b372cc..913c6cca3 100644
--- a/juju/action.py
+++ b/juju/action.py
@@ -2,9 +2,15 @@
 
 
 class Action(model.ModelEntity):
+
+    def __init__(self, entity_id, model, history_index=-1, connected=True):
+        super().__init__(entity_id, model, history_index, connected)
+        self.results = {}
+
     @property
     def status(self):
         return self.data['status']
 
     async def wait(self):
-        return await self.model.get_action_output(self.id)
+        self.results = await self.model.get_action_output(self.id)
+        return self
diff --git a/juju/model.py b/juju/model.py
index 97d521d54..a0e0a2390 100644
--- a/juju/model.py
+++ b/juju/model.py
@@ -2262,7 +2262,19 @@ async def set_constraints(self, constraints):
             constraints=constraints)
 
     async def get_action_output(self, action_uuid, wait=None):
-        """Get the results of an action by ID.
+        """ Get the results of an action by ID.
+
+        :param str action_uuid: Id of the action
+        :param int wait: Time in seconds to wait for action to complete.
+        :return dict: Output from action
+        :raises: :class:`JujuError` if invalid action_uuid
+        """
+        action = await self._get_completed_action(action_uuid, wait=wait)
+        # ActionResult.output is None if the action produced no output
+        return {} if action.output is None else action.output
+
+    async def _get_completed_action(self, action_uuid, wait=None):
+        """Get the completed internal _definitions.Action object.
 
         :param str action_uuid: Id of the action
         :param int wait: Time in seconds to wait for action to complete.
@@ -2288,13 +2300,8 @@ async def _wait_for_action_status():
         await jasyncio.wait_for(
             _wait_for_action_status(),
             timeout=wait)
-        action_output = await action_facade.Actions(entities=entity)
-        # ActionResult.output is None if the action produced no output
-        if action_output.results[0].output is None:
-            output = {}
-        else:
-            output = action_output.results[0].output
-        return output
+        action_results = await action_facade.Actions(entities=entity)
+        return action_results.results[0]
 
     async def get_action_status(self, uuid_or_prefix=None, name=None):
         """Get the status of all actions, filtered by ID, ID prefix, or name.
diff --git a/juju/unit.py b/juju/unit.py
index cecccadac..e65a824b0 100644
--- a/juju/unit.py
+++ b/juju/unit.py
@@ -227,22 +227,27 @@ async def run_action(self, action_name, **params):
 
         """
         action_facade = client.ActionFacade.from_connection(self.connection)
-
         log.debug('Starting action `%s` on %s', action_name, self.name)
 
-        res = await action_facade.EnqueueOperation(actions=[client.Action(
+        old_client = self.connection.is_using_old_client
+
+        op = action_facade.Enqueue if old_client else action_facade.EnqueueOperation
+        res = await op(actions=[client.Action(
             name=action_name,
             parameters=params,
             receiver=self.tag,
         )])
-        action = res.actions[0].result
-        error = res.actions[0].error
+
+        _action = res.results[0] if old_client else res.actions[0]
+        action = _action.action
+        error = _action.error
+
         if error and error.code == 'not found':
             raise ValueError('Action `%s` not found on %s' % (action_name,
                                                               self.name))
         elif error:
             raise Exception('Unknown action error: %s' % error.serialize())
-        action_id = action[len('action-'):]
+        action_id = action.tag[len('action-'):]
         log.debug('Action started as %s', action_id)
         # we mustn't use wait_for_action because that blocks until the
         # action is complete, rather than just being in the model