Looking for design and code correctness feedback. I am writing a package that will provide a light weight wrapper for pymongo insert (and eventually other methods), in order to perform basic validation, logging, etc.
pymongobean/collection.py
import pymongo
import inspect
class MongoBeanWriteWrapper:
def __init__(self, before=None, after=None):
self.before = before
self.after = after
def __call__(self, fn):
def wrapped(*args, **kwargs):
if self.before:
self.before(*args, **kwargs)
result = fn(*args, **kwargs)
if self.after:
self.after(result)
return result
return wrapped
class MongoBeanCollection(pymongo.collection.Collection):
def __init__(self, database, name, create=False, **kwargs):
self._before_insert = None
self._after_insert = None
super(MongoBeanCollection, self).__init__(database, name)
def __getattr__(self, name):
return MongoBeanCollection(self.__database, u"%s.%s" % (self.__name, name))
def before_insert(self, fn):
argspec = inspect.getargspec(fn)
if fn and argspec[0]: # insert must at least have one required arg
self._before_insert = fn
else:
raise InvalidMongoBeanBeforeInsert('before_insert expects a function argument')
def after_insert(self, fn):
argspec = inspect.getargspec(fn)
if fn and argspec[0]:
self._after_insert = fn
else:
raise InvalidMongoBeanAfterInsert('after_insert expects a function argument')
def insert(self, doc_or_docs, manipulate=True,
safe=None, check_keys=True, continue_on_error=False, **kwargs):
@MongoBeanWriteWrapper(self._before_insert, self._after_insert)
def call(doc_or_docs, manipulate, safe, check_keys,
continue_on_error, **kwargs):
return self.base_insert(doc_or_docs, manipulate, safe, check_keys,
continue_on_error, **kwargs)
return call(doc_or_docs, manipulate, safe, check_keys,
continue_on_error, **kwargs)
def base_insert(self, doc_or_docs, manipulate, safe, check_keys,
continue_on_error, **kwargs):
return super(MongoBeanCollection, self).insert(doc_or_docs, manipulate,
safe, check_keys,
continue_on_error,
**kwargs)
class InvalidMongoBeanBeforeInsert(ValueError):
pass
class InvalidMongoBeanAfterInsert(ValueError):
pass
tests
import unittest
import pymongo
import time
import mock
#assumes local install off mongodb is running
#will create a test db if it does not exist
class PyMongoBeanDBTests(unittest.TestCase):
def _get_db(self):
from pymongobean.database import MongoBeanDB
pymongo_conn = pymongo.connection.Connection('mongodb://localhost')
db = pymongo_conn.test
return MongoBeanDB(db)
def test_should_be_instance_of_pymongo_db(self):
mongobean_db = self._get_db()
self.assertTrue(isinstance(mongobean_db, pymongo.database.Database),
str(mongobean_db.__class__))
def test_should_call_validate_function_before_insert_into_user_collection(self):
db = self._get_db()
validate_mock = mock.Mock()
def validate(data, *args, **kwargs):
validate_mock(data) # we just want to make sure before_insert is called with data
user_col = db.user
user_col.before_insert(validate)
data = {'username': timestamp()}
user_col.base_insert = mock.Mock(return_value='testid')
oid = user_col.insert(data)
validate_mock.assert_called_once_with(data)
user_col.base_insert.assert_called_once_with(data,
True, None, True, False)
self.assertTrue(oid, oid)
def test_should_raise_invalid_mongobean_before_insert_error(self):
db = self._get_db()
user_col = db.user
def missing_arg_func():
pass
from pymongobean.collection import InvalidMongoBeanBeforeInsert
with self.assertRaises(InvalidMongoBeanBeforeInsert):
user_col.before_insert(missing_arg_func)
def test_should_raise_invalid_mongobean_after_insert_error(self):
db = self._get_db()
user_col = db.user
def missing_arg_func():
pass
from pymongobean.collection import InvalidMongoBeanAfterInsert
with self.assertRaises(InvalidMongoBeanAfterInsert):
user_col.after_insert(missing_arg_func)
def test_should_call_log_function_after_insert(self):
db = self._get_db()
log_mock = mock.Mock()
def log(data, *args, **kwargs):
log_mock(data)
user_col = db.user
user_col.after_insert(log)
data = {'username': timestamp()}
user_col.base_insert = mock.Mock(return_value='testid')
oid = user_col.insert(data)
log_mock.assert_called_once_with(oid)
user_col.base_insert.assert_called_once_with(data,
True, None, True, False)
self.assertTrue(oid, oid)
def timestamp():
return time.strftime('%Y%m%d%H%M%S')
if __name__ == '__main__':
unittest.main()
pymongo driver for your reference: https://github.com/mongodb/mongo-python-driver/blob/master/pymongo/collection.py