diff --git a/pillar/application/utils/__init__.py b/pillar/application/utils/__init__.py index 42ef8423..4a444375 100644 --- a/pillar/application/utils/__init__.py +++ b/pillar/application/utils/__init__.py @@ -4,9 +4,11 @@ import datetime import functools import logging -import bson +import bson.objectid from eve import RFC1123_DATE_FORMAT from flask import current_app +from werkzeug import exceptions as wz_exceptions + __all__ = ('remove_private_keys', 'PillarJSONEncoder') log = logging.getLogger(__name__) @@ -80,3 +82,20 @@ def project_get_node_type(project_document, node_type_node_name): return next((node_type for node_type in project_document['node_types'] if node_type['name'] == node_type_node_name), None) + + +def str2id(document_id): + """Returns the document ID as ObjectID, or raises a BadRequest exception. + + :type document_id: str + :rtype: bson.ObjectId + :raises: wz_exceptions.BadRequest + """ + + if not document_id: + raise wz_exceptions.BadRequest('Invalid object ID %r', document_id) + + try: + return bson.ObjectId(document_id) + except bson.objectid.InvalidId: + raise wz_exceptions.BadRequest('Invalid object ID %r', document_id) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..e616e148 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,31 @@ +# -*- encoding: utf-8 -*- + +from bson import ObjectId +from werkzeug.exceptions import BadRequest + +from common_test_class import AbstractPillarTest + + +class Str2idTest(AbstractPillarTest): + def test_happy(self): + from application.utils import str2id + + def happy(str_id): + self.assertEqual(ObjectId(str_id), str2id(str_id)) + + happy(24 * 'a') + happy(12 * 'a') + happy(u'577e23ad98377323f74c368c') + + def test_unhappy(self): + from application.utils import str2id + + def unhappy(str_id): + self.assertRaises(BadRequest, str2id, str_id) + + unhappy(13 * 'a') + unhappy(u'577e23ad 8377323f74c368c') + unhappy(u'김치') # Kimchi + unhappy('') + unhappy(u'') + unhappy(None)