Use status code instead of simple message for testingHEADmaster

1 files changed, 16 insertions, 10 deletions

diff --git a/src/lb_app/app_db/user.py b/src/lb_app/app_db/user.py
index 7aeab7c..9f7988d 100644
--- a/src/lb_app/app_db/user.py
+++ b/src/lb_app/app_db/user.py
@@ -14,29 +14,35 @@ class User(object):
     def register(self, username):
         db = MongoDB(self.app)
         mongo = db.connection()
-        error = None
+        error = {}
 
         if not username:
-            error = 'Username is required.'
+            error = {'msg': 'Username is required.', 'code': 400}
         elif len(list(mongo.db.users.find({"username": username}))) > 0:
-            error = 'User {} is already registered.'.format(username)
+            error = {
+                'msg': 'User {} is already registered.'.format(username),
+                'code': 409
+            }
 
-        if error is None:
+        if not error:
             mongo.db.users.insert({"username": username})
-            return None
+            return {'msg': 'User {} added.'.format(username), 'code': 201}
         return error
 
     def remove(self, username):
         db = MongoDB(self.app)
         mongo = db.connection()
-        error = None
+        error = {}
 
         if not username:
-            error = 'Username is required.'
+            error = {'msg': 'Username is required', 'code': 400}
         elif len(list(mongo.db.users.find({"username": username}))) == 0:
-            error = "User {} not exist. So it's good.".format(username)
+            error = {
+                'msg': "User {} not exist. So it's good.".format(username),
+                'code': 404
+            }
 
-        if error is None:
+        if not error:
             mongo.db.users.remove({"username": username})
-            return None
+            return {'msg': 'User {} deleted.'.format(username), 'code': 200}
         return error