Add working test setup

11 files changed, 167 insertions, 144 deletions

diff --git a/README.md b/README.md
index fbf6f913..cd8bd956 100644
--- a/README.md
+++ b/README.md
@@ -121,8 +121,4 @@ To try a different query, use the Postman `Builder` to edit the method, path, bo
 
 When editing the web server code, restart the server (`CTRL` + `c` followed by `python main.py config.json` in the console running the server) to see the result of your changes.
 
-To format all files, run the following command:
-
-```bash
-yapf **/*.py -i
-```
+To format all files, run `format.sh` in this directory.
diff --git a/conftest.py b/conftest.py
new file mode 100644
index 00000000..1e404c1c
--- /dev/null
+++ b/conftest.py
@@ -0,0 +1,13 @@
+import pytest
+
+from main import FLASK_CORE_APP
+
+
+@pytest.fixture
+def client(mocker):
+    """Returns a Flask API client to interact with."""
+    FLASK_CORE_APP.config['TESTING'] = True
+    mocker.patch('opendc.util.database.DB')
+
+    with FLASK_CORE_APP.test_client() as client:
+        yield client
diff --git a/format.sh b/format.sh
new file mode 100644
index 00000000..18cba452
--- /dev/null
+++ b/format.sh
@@ -0,0 +1 @@
+yapf **/*.py -i
diff --git a/main.py b/main.py
index 6f8ddbef..70b77cae 100644
--- a/main.py
+++ b/main.py
@@ -10,33 +10,29 @@ from oauth2client import client, crypt
 from flask_cors import CORS
 
 from opendc.models.user import User
-from opendc.util import exceptions, rest, path_parser, database
+from opendc.util import rest, path_parser, database
+from opendc.util.exceptions import AuthorizationTokenError, RequestInitializationError
 
-if len(sys.argv) < 2:
-    print("config file path not given as argument")
-    sys.exit(1)
+TEST_MODE = "OPENDC_FLASK_TESTING" in os.environ
 
-# Get keys from config file
-with open(sys.argv[1]) as f:
-    KEYS = json.load(f)
-
-STATIC_ROOT = os.path.join(KEYS['ROOT_DIR'], 'opendc-frontend', 'build')
-
-database.init_connection_pool(user=KEYS['OPENDC_DB_USERNAME'],
-                              password=KEYS['OPENDC_DB_PASSWORD'],
-                              database=KEYS['OPENDC_DB'],
-                              host='localhost',
-                              port=27017)
+if TEST_MODE:
+    STATIC_ROOT = os.curdir
+else:
+    database.initialize_database(user=os.environ['OPENDC_DB_USERNAME'],
+                                 password=os.environ['OPENDC_DB_PASSWORD'],
+                                 database=os.environ['OPENDC_DB'],
+                                 host='localhost')
+    STATIC_ROOT = os.path.join(os.environ['OPENDC_ROOT_DIR'], 'opendc-frontend', 'build')
 
 FLASK_CORE_APP = Flask(__name__, static_url_path='', static_folder=STATIC_ROOT)
-FLASK_CORE_APP.config['SECRET_KEY'] = KEYS['FLASK_SECRET']
-if 'localhost' in KEYS['SERVER_BASE_URL']:
+FLASK_CORE_APP.config['SECRET_KEY'] = os.environ['OPENDC_FLASK_SECRET']
+if 'localhost' in os.environ['OPENDC_SERVER_BASE_URL']:
     CORS(FLASK_CORE_APP)
 
 compress = Compress()
 compress.init_app(FLASK_CORE_APP)
 
-if 'localhost' in KEYS['SERVER_BASE_URL']:
+if 'OPENDC_SERVER_BASE_URL' in os.environ or 'localhost' in os.environ['OPENDC_SERVER_BASE_URL']:
     SOCKET_IO_CORE = flask_socketio.SocketIO(FLASK_CORE_APP, cors_allowed_origins="*")
 else:
     SOCKET_IO_CORE = flask_socketio.SocketIO(FLASK_CORE_APP)
@@ -57,9 +53,9 @@ def sign_in():
         return 'No idtoken provided', 401
 
     try:
-        idinfo = client.verify_id_token(token, KEYS['OAUTH_CLIENT_ID'])
+        idinfo = client.verify_id_token(token, os.environ['OPENDC_OAUTH_CLIENT_ID'])
 
-        if idinfo['aud'] != KEYS['OAUTH_CLIENT_ID']:
+        if idinfo['aud'] != os.environ['OPENDC_OAUTH_CLIENT_ID']:
             raise crypt.AppIdentityError('Unrecognized client.')
 
         if idinfo['iss'] not in ['accounts.google.com', 'https://accounts.google.com']:
@@ -127,8 +123,7 @@ def api_call(version, endpoint_path):
 @FLASK_CORE_APP.route('/my-auth-token')
 def serve_web_server_test():
     """Serve the web server test."""
-
-    return send_from_directory(os.path.join(KEYS['ROOT_DIR'], 'opendc-web-server', 'static'), 'index.html')
+    return send_from_directory(STATIC_ROOT, 'index.html')
 
 
 @FLASK_CORE_APP.route('/')
@@ -144,50 +139,48 @@ def serve_index(simulation_id=None, experiment_id=None):
 @SOCKET_IO_CORE.on('request')
 def receive_message(message):
     """"Receive a SocketIO request"""
+    (req, res) = _process_message(message)
 
-    (request, response) = _process_message(message)
-
-    print(
-        f'Socket:\t{request.method} to `/{request.path}` resulted in {response.status["code"]}: {response.status["description"]}'
-    )
+    print(f'Socket:\t{req.method} to `/{req.path}` resulted in {res.status["code"]}: {res.status["description"]}')
     sys.stdout.flush()
 
-    flask_socketio.emit('response', response.to_JSON(), json=True)
+    flask_socketio.emit('res', res.to_JSON(), json=True)
 
 
 def _process_message(message):
     """Process a request message and return the response."""
 
     try:
-        request = rest.Request(message)
-        response = request.process()
+        req = rest.Request(message)
+        res = req.process()
 
-        return (request, response)
+        return req, res
 
-    except exceptions.AuthorizationTokenError as e:
-        response = rest.Response(401, 'Authorization error')
-        response.id = message['id']
+    except AuthorizationTokenError:
+        res = rest.Response(401, 'Authorization error')
+        res.id = message['id']
 
-    except exceptions.RequestInitializationError as e:
-        response = rest.Response(400, str(e))
-        response.id = message['id']
+    except RequestInitializationError as e:
+        res = rest.Response(400, str(e))
+        res.id = message['id']
 
         if not 'method' in message:
             message['method'] = 'UNSPECIFIED'
         if not 'path' in message:
             message['path'] = 'UNSPECIFIED'
 
-    except Exception as e:
-        response = rest.Response(500, 'Internal server error')
+    except Exception:
+        res = rest.Response(500, 'Internal server error')
         if 'id' in message:
-            response.id = message['id']
+            res.id = message['id']
         traceback.print_exc()
 
-    request = rest.Request()
-    request.method = message['method']
-    request.path = message['path']
+    req = rest.Request()
+    req.method = message['method']
+    req.path = message['path']
 
-    return (request, response)
+    return req, res
 
 
-SOCKET_IO_CORE.run(FLASK_CORE_APP, host='0.0.0.0', port=8081)
+if __name__ == '__main__':
+    SOCKET_IO_CORE.run(FLASK_CORE_APP, host='0.0.0.0', port=8081)
diff --git a/opendc/api/v2/tiles/tileId/rack/machines/position/endpoint.py b/opendc/api/v2/tiles/tileId/rack/machines/position/endpoint.py
index e769a049..72fd44d5 100644
--- a/opendc/api/v2/tiles/tileId/rack/machines/position/endpoint.py
+++ b/opendc/api/v2/tiles/tileId/rack/machines/position/endpoint.py
@@ -96,7 +96,7 @@ def PUT(request):
         return Response(409, 'Rack position occupied.')
 
     except Exception as e:
-        print e
+        print(e)
         return Response(400, 'Invalid Machine.')
 
     # Return this Machine
diff --git a/opendc/api/v2/users/endpoint.py b/opendc/api/v2/users/endpoint.py
index 2712c625..dca509ed 100644
--- a/opendc/api/v2/users/endpoint.py
+++ b/opendc/api/v2/users/endpoint.py
@@ -2,7 +2,7 @@ from werkzeug.exceptions import abort
 
 from opendc.models.user import User
 from opendc.util import exceptions
-from opendc.util.database import fetch_one, insert
+from opendc.util.database import DB
 from opendc.util.rest import Response
 
 
@@ -14,7 +14,7 @@ def GET(request):
     except exceptions.ParameterError as e:
         return Response(400, str(e))
 
-    user = fetch_one({'email': request.params_query['email']}, 'users')
+    user = DB.fetch_one({'email': request.params_query['email']}, 'users')
 
     if user is not None:
         return Response(404, f'User with email {request.params_query["email"]} not found')
@@ -32,7 +32,7 @@ def POST(request):
 
     request.params_body['user']['googleId'] = request.google_id
     user = request.params_body['user']
-    existing_user = fetch_one({'googleId': user['googleId']}, 'users')
+    existing_user = DB.fetch_one({'googleId': user['googleId']}, 'users')
 
     if existing_user is not None:
         return Response(409, '{} already exists.'.format(existing_user))
diff --git a/opendc/api/v2/users/test_endpoint.py b/opendc/api/v2/users/test_endpoint.py
new file mode 100644
index 00000000..ffe2ce02
--- /dev/null
+++ b/opendc/api/v2/users/test_endpoint.py
@@ -0,0 +1,2 @@
+def test_get_user_missing_parameter(client):
+    print(client.get('/api/v2/users'))
diff --git a/opendc/util/database.py b/opendc/util/database.py
index 0e424aa4..24572279 100644
--- a/opendc/util/database.py
+++ b/opendc/util/database.py
@@ -7,91 +7,87 @@ from pymongo import MongoClient
 
 DATETIME_STRING_FORMAT = '%Y-%m-%dT%H:%M:%S'
 CONNECTION_POOL = None
+DB = None
 
 
-def init_connection_pool(user, password, database, host, port):
-    user = urllib.parse.quote_plus(user)  # TODO: replace this with environment variable
-    password = urllib.parse.quote_plus(password)  # TODO: same as above
-    database = urllib.parse.quote_plus(database)
-    host = urllib.parse.quote_plus(host)
+class Database:
+    def __init__(self, user, password, database, host):
+        user = urllib.parse.quote_plus(user)  # TODO: replace this with environment variable
+        password = urllib.parse.quote_plus(password)  # TODO: same as above
+        database = urllib.parse.quote_plus(database)
+        host = urllib.parse.quote_plus(host)
 
-    global opendcdb
+        client = MongoClient('mongodb://%s:%s@%s/default_db?authSource=%s' % (user, password, host, database))
+        self.opendc_db = client.opendc
 
-    client = MongoClient('mongodb://%s:%s@%s/default_db?authSource=%s' % (user, password, host, database))
-    opendcdb = client.opendc
+    def fetch_one(self, query, collection):
+        """Uses existing mongo connection to return a single (the first) document in a collection matching the given
+        query as a JSON object.
 
+        The query needs to be in json format, i.e.: `{'name': prefab_name}`.
+        """
+        bson = getattr(self.opendc_db, collection).find_one(query)
 
-def fetch_one(query, collection):
-    """Uses existing mongo connection to return a single (the first) document in a collection matching the given
-    query as a JSON object.
+        return self.convert_bson_to_json(bson)
 
-    The query needs to be in json format, i.e.: `{'name': prefab_name}`.
-    """
-    bson = getattr(opendcdb, collection).find_one(query)
+    def fetch_all(self, query, collection):
+        """Uses existing mongo connection to return all documents matching a given query, as a list of JSON objects.
 
-    return convert_bson_to_json(bson)
+        The query needs to be in json format, i.e.: `{'name': prefab_name}`.
+        """
+        results = []
+        cursor = getattr(self.opendc_db, collection).find(query)
+        for doc in cursor:
+            results.append(self.convert_bson_to_json(doc))
+        return results
 
+    def insert(self, obj, collection):
+        """Updates an existing object."""
+        bson = getattr(self.opendc_db, collection).insert(obj)
 
-def fetch_all(query, collection):
-    """Uses existing mongo connection to return all documents matching a given query, as a list of JSON objects.
+        return self.convert_bson_to_json(bson)
 
-    The query needs to be in json format, i.e.: `{'name': prefab_name}`.
-    """
-    results = []
-    cursor = getattr(opendcdb, collection).find(query)
-    for doc in cursor:
-        results.append(convert_bson_to_json(doc))
-    return results
+    def update(self, _id, obj, collection):
+        """Updates an existing object."""
+        bson = getattr(self.opendc_db, collection).update({'_id': _id}, obj)
 
+        return self.convert_bson_to_json(bson)
 
-def insert(obj, collection):
-    """Updates an existing object."""
-    bson = getattr(opendcdb, collection).insert(obj)
+    def delete_one(self, query, collection):
+        """Deletes one object matching the given query.
 
-    return convert_bson_to_json(bson)
+        The query needs to be in json format, i.e.: `{'name': prefab_name}`.
+        """
+        bson = getattr(self.opendc_db, collection).delete_one(query)
 
+        return self.convert_bson_to_json(bson)
 
-def update(_id, obj, collection):
-    """Updates an existing object."""
-    bson = getattr(opendcdb, collection).update({'_id': _id}, obj)
+    def delete_all(self, query, collection):
+        """Deletes all objects matching the given query.
 
-    return convert_bson_to_json(bson)
+        The query needs to be in json format, i.e.: `{'name': prefab_name}`.
+        """
+        bson = getattr(self.opendc_db, collection).delete_many(query)
 
+        return self.convert_bson_to_json(bson)
 
-def delete_one(query, collection):
-    """Deletes one object matching the given query.
+    @staticmethod
+    def convert_bson_to_json(bson):
+        """Converts a BSON representation to JSON and returns the JSON representation."""
+        json_string = dumps(bson)
+        return json.loads(json_string)
 
-    The query needs to be in json format, i.e.: `{'name': prefab_name}`.
-    """
-    bson = getattr(opendcdb, collection).delete_one(query)
+    @staticmethod
+    def datetime_to_string(datetime_to_convert):
+        """Return a database-compatible string representation of the given datetime object."""
+        return datetime_to_convert.strftime(DATETIME_STRING_FORMAT)
 
-    return convert_bson_to_json(bson)
+    @staticmethod
+    def string_to_datetime(string_to_convert):
+        """Return a datetime corresponding to the given string representation."""
+        return datetime.strptime(string_to_convert, DATETIME_STRING_FORMAT)
 
 
-def delete_all(query, collection):
-    """Deletes all objects matching the given query.
-
-    The query needs to be in json format, i.e.: `{'name': prefab_name}`.
-    """
-    bson = getattr(opendcdb, collection).delete_many(query)
-
-    return convert_bson_to_json(bson)
-
-
-def convert_bson_to_json(bson):
-    # Convert BSON representation to JSON
-    json_string = dumps(bson)
-    # Load as a JSON object
-    return json.loads(json_string)
-
-
-def datetime_to_string(datetime_to_convert):
-    """Return a database-compatible string representation of the given datetime object."""
-
-    return datetime_to_convert.strftime(DATETIME_STRING_FORMAT)
-
-
-def string_to_datetime(string_to_convert):
-    """Return a datetime corresponding to the given string representation."""
-
-    return datetime.strptime(string_to_convert, DATETIME_STRING_FORMAT)
+def initialize_database(user, password, database, host):
+    global DB
+    DB = Database(user, password, database, host)
diff --git a/opendc/util/rest.py b/opendc/util/rest.py
index d892a358..2312b199 100644
--- a/opendc/util/rest.py
+++ b/opendc/util/rest.py
@@ -1,14 +1,12 @@
 import importlib
 import json
+import os
 import sys
 
 from oauth2client import client, crypt
 
 from opendc.util import exceptions, parameter_checker
 
-with open(sys.argv[1]) as f:
-    KEYS = json.load(f)
-
 
 class Request(object):
     """WebSocket message to REST request mapping."""
@@ -45,6 +43,7 @@ class Request(object):
             module_base = 'opendc.api.{}.endpoint'
             module_path = self.path.replace('/', '.')
 
+            print(module_base.format(module_path))
             self.module = importlib.import_module(module_base.format(module_path))
         except ImportError:
             raise exceptions.UnimplementedEndpointError('Unimplemented endpoint: {}.'.format(self.path))
@@ -60,31 +59,14 @@ class Request(object):
 
         # Verify the user
 
+        if "OPENDC_FLASK_TESTING" in os.environ:
+            return
+
         try:
             self.google_id = self._verify_token(self.token)
-
         except crypt.AppIdentityError as e:
             raise exceptions.AuthorizationTokenError(e)
 
-    def _verify_token(self, token):
-        """Return the ID of the signed-in user.
-        
-        Or throw an Exception if the token is invalid.
-        """
-
-        try:
-            idinfo = client.verify_id_token(token, KEYS['OAUTH_CLIENT_ID'])
-        except Exception as e:
-            raise crypt.AppIdentityError('Exception caught trying to verify ID token: {}'.format(e))
-
-        if idinfo['aud'] != KEYS['OAUTH_CLIENT_ID']:
-            raise crypt.AppIdentityError('Unrecognized client.')
-
-        if idinfo['iss'] not in ['accounts.google.com', 'https://accounts.google.com']:
-            raise crypt.AppIdentityError('Wrong issuer.')
-
-        return idinfo['sub']
-
     def check_required_parameters(self, **kwargs):
         """Raise an error if a parameter is missing or of the wrong type."""
 
@@ -108,6 +90,27 @@ class Request(object):
 
         return json.dumps(self.message)
 
+    @staticmethod
+    def _verify_token(token):
+        """Return the ID of the signed-in user.
+
+        Or throw an Exception if the token is invalid.
+        """
+
+        try:
+            id_info = client.verify_id_token(token, os.environ['OPENDC_OAUTH_CLIENT_ID'])
+        except Exception as e:
+            print(e)
+            raise crypt.AppIdentityError('Exception caught trying to verify ID token: {}'.format(e))
+
+        if id_info['aud'] != os.environ['OPENDC_OAUTH_CLIENT_ID']:
+            raise crypt.AppIdentityError('Unrecognized client.')
+
+        if id_info['iss'] not in ['accounts.google.com', 'https://accounts.google.com']:
+            raise crypt.AppIdentityError('Wrong issuer.')
+
+        return id_info['sub']
+
 
 class Response(object):
     """Response to websocket mapping"""
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 00000000..775a8ff4
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,5 @@
+[pytest]
+env =
+    OPENDC_FLASK_TESTING=True
+    OPENDC_FLASK_SECRET=Secret
+    OPENDC_SERVER_BASE_URL=localhost
diff --git a/setup.py b/setup.py
index 23152cc4..7fdb70f1 100644
--- a/setup.py
+++ b/setup.py
@@ -24,8 +24,22 @@ setup(
     ],
     keywords='opendc datacenter simulation web-server',
     packages=['opendc'],
+    # yapf: disable
     install_requires=[
-        'flask==1.0.2', 'flask-socketio==3.0.2', 'oauth2client==4.1.3', 'eventlet==0.24.1', 'flask-compress==1.4.0',
-        'flask-cors==3.0.8', 'pyasn1-modules==0.2.2', 'six==1.11.0', 'pymongo==3.10.1', 'bson==0.5.10', 'yapf==0.30.0'
+        'flask==1.0.2',
+        'flask-socketio==3.0.2',
+        'oauth2client==4.1.3',
+        'eventlet==0.24.1',
+        'flask-compress==1.4.0',
+        'flask-cors==3.0.8',
+        'pyasn1-modules==0.2.2',
+        'six==1.11.0',
+        'pymongo==3.10.1',
+        'bson==0.5.10',
+        'yapf==0.30.0',
+        'pytest==5.4.3',
+        'pytest-mock==3.1.1',
+        'pytest-env==0.6.2',
     ],
+    # yapf: enable
 )