linkyndy
linkyndy

Reputation: 17900

Mocking a database connection in a Model class in Python

I have written a Model class which wraps some RethinkDB functionality in an object. Part of it looks like:

class Model(object):
    __metaclass__ = ModelType

    def __init__(self, **kwargs):
        self._fields = kwargs
        self._saved = False

    def save(self):
        if self._saved:
            id = self._fields.pop('id')
            result = r.table(self._tablename)
                      .get(id)
                      .update(self._fields, return_vals=True)
                      .run(self.connection())
        else:
            result = r.table(self._tablename)
                      .insert(self._fields, return_vals=True)
                      .run(self.connection())

        if result['errors'] > 0:
            raise ActionError(result['first_error'])

        self._fields = result['new_val']
        self._saved = True

Now, I would like to test whether the save() method does what it is supposed to do. Since I am using a database connection for this operation, creating and removing a connection before the test is expensive and useless, since I only need to test that the correct methods on the r object are called. Obviously, mocking is the solution to this, but I am facing some difficulties in what and how should I mock the r object. I admit I am new to mocking concepts.

So, in my method that tests save(), how can I test whether the update()/insert() have been called with the right arguments, and also, how can I assign some return values to them?

Upvotes: 2

Views: 2072

Answers (1)

Chris Clarke
Chris Clarke

Reputation: 2181

so, you need to mock r:

r = mock.Mock()
r.table = mock.Mock(return_value=r)
r.insert = mock.Mock(return_value=r)

Model.save()

r.insert.assert_called_with("args", "you're" "expecting")

I'm assuming you can extrapolate from there

Upvotes: 3

Related Questions