Use status code instead of simple message for testingHEADmaster

2 files changed, 23 insertions, 17 deletions

diff --git a/src/lb_app/api.py b/src/lb_app/api.py
index 9827828..9ac2bef 100644
--- a/src/lb_app/api.py
+++ b/src/lb_app/api.py
@@ -27,7 +27,11 @@ user_model = api.model('User', {
 class Users(Resource):
     def get(self):
         """ Get user list """
-        return loads(dumps(user.all()))
+        users = loads(dumps(user.all()))
+        if not users:
+            return loads('{"msg": "No users"}'), 404
+            print(users)
+        return loads(dumps(user.all())), 200
 
 
 @api.route('/user/register')
@@ -38,9 +42,7 @@ class Register(Resource):
         args = request.get_json(force=True)
         status = user.register(args.get('username'))
 
-        if status is None:
-            return loads('{"msg": "done"}'), 201
-        return loads('{"msg": "'+status+'"}'), 200
+        return loads('{"msg": "'+status['msg']+'"}'), status['code']
 
 @api.route('/users/<username>')
 @api.doc(params={'username': 'Jean'})
@@ -49,6 +51,4 @@ class Remove(Resource):
       """ Remove a user """
       status = user.remove(username)
 
-      if status is None:
-          return loads('{"msg": "done"}'), 201
-      return loads('{"msg": "'+status+'"}'), 200
+      return loads('{"msg": "'+status['msg']+'"}'), status['code']
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