Take the 2-minute tour ×
Code Review Stack Exchange is a question and answer site for peer programmer code reviews. It's 100% free, no registration required.

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

share|improve this question
add comment

Know someone who can answer? Share a link to this question via email, Google+, Twitter, or Facebook.

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Browse other questions tagged or ask your own question.