Flask 单元测试——模拟 Aerospike DB

Question

I have the following endpoint,

    @developer_blueprint.route("/init_db", methods=["POST"])
    def initialize_database():
       try:
          upload_data(current_app)
          logger.debug("Database entries upload.")
          return jsonify({"result": "Database entries uploaded."}), 201
       except Exception as e:
          return jsonify({"error": str(e)})

    def upload_data(app):
       with open("src/core/data/data.json") as data_file:
          data = json.load(data_file)
          try:
             current_app.db.put(("somenamespace", "test", "default"), data, None)
          except Exception as e:
             raise e

I'm trying to figure out how to unit test this (we need to get coverage on our code). Do I just mock up app.db? How can I do that?

Any suggestions would be appreciated.

Answer 1

It is not uncommon to mock database calls for unit testing using something like unittest.mock and then run Aerospike in a container or VM for end-to-end testing.

However, keep in mind that the Aerospike Python client library is written in C for better performance and thus it is not easy to do partial patching (aka "monkey patching"). For example, you will get a TypeError: can't set attributes of built-in/extension type if you try to simply patch out aerospike.Client.put .

One approach is to create a mock client object to replace or sub-class the Aerospike client object. The implementation of this mock object depends on your code and the cases you are testing for.

Take the following example code in which app.db is an instance of the Aerospike client library:

# example.py
import aerospike
import json

class App(object):
    db = None
    def __init__(self):
        config = {'hosts': [('127.0.0.1', 3000)]}
        self.db = aerospike.client(config).connect()

def upload_data(app):
    with open("data.json") as data_file:
        data = json.load(data_file)

    try:
        app.db.put(("ns1", "test", "default"), data, None)
    except Exception as e:
            raise e

if __name__ == "__main__":
    app = App()
    upload_data(app)

In writing unit tests for the upload_data function let's assume you want to test for a success case which is determined to mean that the put method is called and no exceptions are raised:

# test.py
from unittest import TestCase, main
from unittest.mock import PropertyMock, patch
from example import App, upload_data
from aerospike import Client, exception


class MockClient(Client):
    def __init__(self, *args, **kwargs):
        pass

    def put(self, *args, **kwargs):
        return 0


class ExampleTestCase(TestCase):
    def test_upload_data_success(self):
        with patch.object(App, 'db', new_callable=PropertyMock) as db_mock:
            db_mock.return_value = client = MockClient()
            app = App()
            with patch.object(client, 'put') as put_mock:
                upload_data(app)
                put_mock.assert_called()


if __name__ == '__main__':
    main()

In the test_upload_data_success method the App.db property is patched with the MockClient class instead of the aerospike.Client class. The put method of the MockClient instance is also patched so that it can be asserted that the put method gets called after upload_data is called.

To test that an exception raised by the Aerospike client is re-raised from the upload_data function, the MockClient class can be modified to raise an exception explicitly:

# test.py
from unittest import TestCase, main
from unittest.mock import PropertyMock, patch
from example import App, upload_data
from aerospike import Client, exception


class MockClient(Client):
    def __init__(self, *args, **kwargs):
        self.put_err = None
        if 'put_err' in kwargs:
            self.put_err = kwargs['put_err']

    def put(self, *args, **kwargs):
        if self.put_err:
            raise self.put_err
        else:
            return 0


class ExampleTestCase(TestCase):
    def test_upload_data_success(self):
        with patch.object(App, 'db', new_callable=PropertyMock) as db_mock:
            db_mock.return_value = client = MockClient()
            app = App()
            with patch.object(client, 'put') as put_mock:
                upload_data(app)
                put_mock.assert_called()

    def test_upload_data_error(self):
        with patch.object(App, 'db', new_callable=PropertyMock) as db_mock:
            db_mock.return_value = MockClient(put_err=exception.AerospikeError)
            app = App()
            with self.assertRaises(exception.AerospikeError):
                upload_data(app)


if __name__ == '__main__':
    main()