tikej
tikej

Reputation: 312

How to test if SQL Alchemy transaction was called?

How do I test following code, where engine is originally meant to be SQLAlchemy engine object? The test case I provided unfortunately doesn't work. I was able to check if the begin method of engine was called, but not the execute statement. The second execute_mock is never called after executing fetch_data method.

class A():
    def __init__(self, engine):
        self.engine = engine

    def fetch_data(self):
        with self.engine.begin() as trans:
            trans.execute("SELECT * FROM XXX")

from unittest.mock import MagicMock
def test_A():
    execute_mock = MagicMock()
    engine_mock = MagicMock()
    engine_mock.begin.return_value.execute = execute_mock

    A(engine_mock)
    execute_mock.assert_not_called()
    A.fetch_data()
    execute_mock.assert_called_with("SELECT * FROM XXX")

Upvotes: 0

Views: 770

Answers (1)

Lin Du
Lin Du

Reputation: 102207

You didn't mock the trans context manager correctly. Here is the unit test solution:

a.py:

class A():
    def __init__(self, engine):
        self.engine = engine

    def fetch_data(self):
        with self.engine.begin() as trans:
            rval = trans.execute("SELECT * FROM XXX")

test_a.py:

import unittest
from unittest.mock import MagicMock, mock_open
from a import A


class TestA(unittest.TestCase):
    def test_fetch_data(self):
        engine_mock = MagicMock()
        trans = engine_mock.begin.return_value.__enter__.return_value
        trans.execute.return_value = 'fake data'
        a = A(engine_mock)
        a.fetch_data()
        engine_mock.begin.assert_called_once()
        trans.execute.assert_called_with("SELECT * FROM XXX")


if __name__ == '__main__':
    unittest.main()

unit test results with 100% coverage:

.
----------------------------------------------------------------------
Ran 1 test in 0.002s

OK
Name                                   Stmts   Miss  Cover   Missing
--------------------------------------------------------------------
src/stackoverflow/61224956/a.py            6      0   100%
src/stackoverflow/61224956/test_a.py      14      0   100%
--------------------------------------------------------------------
TOTAL                                     20      0   100%

python version: Python 3.7.5

Upvotes: 1

Related Questions