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

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.