added exceptions for clients, created test for WS client

1 files changed, 24 insertions, 8 deletions

diff --git a/module/remote/WSClient.py b/module/remote/WSClient.py
index c06bab661..23b5fc3ca 100644
--- a/module/remote/WSClient.py
+++ b/module/remote/WSClient.py
@@ -1,8 +1,11 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-from json_converter import loads, dumps
 from websocket import create_connection
+from httplib import UNAUTHORIZED, FORBIDDEN
+
+from json_converter import loads, dumps
+from ttypes import Unauthorized, Forbidden
 
 class WSClient:
     URL = "ws://localhost:7227/api"
@@ -11,21 +14,34 @@ class WSClient:
         self.url = url or self.URL
         self.ws =  None
 
-    def login(self, username, password):
+    def connect(self):
         self.ws = create_connection(self.URL)
-        return self.call("login", username, password)
 
-    def logout(self):
-        self.call("logout")
+    def close(self):
         self.ws.close()
 
+    def login(self, username, password):
+        if not self.ws: self.connect()
+        return self.call("login", username, password)
+
     def call(self, func, *args, **kwargs):
-        self.ws.send(dumps([func, args, kwargs]))
+        if not self.ws:
+            raise Exception("Not Connected")
+
+        if kwargs:
+            self.ws.send(dumps([func, args, kwargs]))
+        else: # omit kwargs
+            self.ws.send(dumps([func, args]))
+
         code, result = loads(self.ws.recv())
         if code == 404:
             raise AttributeError("Unknown Method")
-        elif code == 505:
-            raise Exception("Remote Exception")
+        elif code == 500:
+            raise Exception("Remote Exception: %s" % result)
+        elif code == UNAUTHORIZED:
+            raise Unauthorized()
+        elif code == FORBIDDEN:
+            raise Forbidden()
 
         return result