Hook up Unit.run() to api

[osm/N2VC.git] / juju / client / connection.py

diff --git a/juju/client/connection.py b/juju/client/connection.py
index 9a631ec..56f9e18 100644 (file)
--- a/juju/client/connection.py
+++ b/juju/client/connection.py
@@ -28,7 +28,13 @@ class Connection:
         client = await Connection.connect_current()
 
     """
-    def __init__(self):
+    def __init__(self, endpoint, uuid, username, password, cacert=None):
+        self.endpoint = endpoint
+        self.uuid = uuid
+        self.username = username
+        self.password = password
+        self.cacert = cacert
+
         self.__request_id__ = 0
         self.addr = None
         self.ws = None
@@ -55,14 +61,14 @@ class Connection:
             result = json.loads(result)
         return result
 
-    async def rpc(self, msg):
+    async def rpc(self, msg, encoder=None):
         self.__request_id__ += 1
         msg['RequestId'] = self.__request_id__
         if'Params' not in msg:
             msg['Params'] = {}
         if "Version" not in msg:
             msg['Version'] = self.facades[msg['Type']]
-        outgoing = json.dumps(msg, indent=2)
+        outgoing = json.dumps(msg, indent=2, cls=encoder)
         await self.ws.send(outgoing)
         result = await self.recv()
         log.debug("send %s got %s", msg, result)
@@ -70,14 +76,28 @@ class Connection:
             raise RuntimeError(result)
         return result
 
+    async def clone(self):
+        """Return a new Connection, connected to the same websocket endpoint
+        as this one.
+
+        """
+        return await Connection.connect(
+            self.endpoint,
+            self.uuid,
+            self.username,
+            self.password,
+            self.cacert,
+        )
+
     @classmethod
     async def connect(cls, endpoint, uuid, username, password, cacert=None):
         url = "wss://{}/model/{}/api".format(endpoint, uuid)
-        client = cls()
+        client = cls(endpoint, uuid, username, password, cacert)
         await client.open(url, cacert)
         server_info = await client.login(username, password)
         client.build_facades(server_info['facades'])
         log.info("Driver connected to juju %s", endpoint)
+
         return client
 
     @classmethod