Use status code instead of simple message for testingHEADmaster

4 files changed, 32 insertions, 26 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
diff --git a/src/tests/test_api.py b/src/tests/test_api.py
index 063f32f..741780e 100644
--- a/src/tests/test_api.py
+++ b/src/tests/test_api.py
@@ -11,12 +11,12 @@ def test_api_get_all_users(client, app):
     assert response.is_json
     assert re.search('username', str(response.data))
 
-@pytest.mark.parametrize(('username', 'message'), (
-    ('', b'{"msg": "Username is required."}'),
-    ('neo', b'{"msg": "User neo is already registered."}'),
-    ('neodarz', b'{"msg": "done"}')
+@pytest.mark.parametrize(('username', 'code'), (
+    ('', 400),
+    ('neo', 409),
+    ('neodarz', 201)
 ))
-def test_api_register(client, app, username, message):
+def test_api_register(client, app, username, code):
     response = client.post('/api/user/register', data='{"username": "'+username+'"}')
     print(response.data)
-    assert message in response.data
+    assert code == response.status_code
diff --git a/src/tests/test_user.py b/src/tests/test_user.py
index f39279e..c6ddbd4 100644
--- a/src/tests/test_user.py
+++ b/src/tests/test_user.py
@@ -12,10 +12,10 @@ def test_get_all_users(app):
 def test_register_user(app):
     with app.app_context():
         user = User(app)
-        assert user.register("neodarz") is None
+        assert user.register("neodarz")['code'] is 201
 
 def test_register_and_remove_user(app):
     with app.app_context():
         user = User(app)
-        assert user.register("neodarz") is None
-        assert user.remove("neodarz") is None
+        assert user.register("neodarz")['code'] is 201
+        assert user.remove("neodarz")['code'] is 200