CODEWITHSUNDEEP
pythonpytestpython-unittestpytest-mock
shikida
shikida

Reputation: 505

How to get the internal state of a mock object using python

I want to use a mock object to emulate the database on a insert operation.

For example, let's suppose I have a method like insert(1) that calls a db connection object (let's call it db_obj) and performs some insert into mytable (col1) values (%s) where %s is 1.

What I had in mind: create some mock object for db_obj that stores the value of col1, so when db_obj.insert(1) is called, the mocked db_obj stores this col1=1 and, then, I can just get the mocked object col1 value and assert that it's 1 as expected.

Does this approach makes sense? If so, how can I do this using pytest?

Here's an example of what I'm trying

from hub.scripts.tasks.notify_job_module import http_success
from hub.scripts.tasks.notify_job_module import proceed_to
from hub.core.connector.mysql_db import MySql
from unittest.mock import patch
import logging

def proceed_to(conn,a,b,c,d,e):
  conn.run_update('insert into table_abc (a,b,c,d,e) values (%s,%s,%s,%s,%s)',[a,b,c,d,e])

class MySql:
# this is the real conn implementation

    def connect(self):
      # ... create the database connection here

    def run_update(self, query, params=None):
      # ... perform some insert into the database here

class TestMySql:
# this is the mock implementation I want to get rid of

    def connect(self):
      pass

    def run_update(self, query, params=None):
      self.result = params

    def get_result(self):
      return self.result

def test_proceed_to():

    logger = logging.getLogger("")
    conn = TestMySql() ## originally, my code uses MySql() here
    conn.connect()
    proceed_to(conn,1,'2',3,4,5)
    assert conn.get_result()[1] == 4

Please notice that I had to replace MySql() with TestMySql(), so what I've done was to implement my own Mock object manually.

This way it works, but I feel like it's obviously not the best approach here. Why?

Because we're talking about mock objects, the definition of proceed_to is irrelevant here :-) the thing is: I had to implement TestMySql.get_result() and store the data I want in self.result to get the result I want, while MySql itself does not has a get_result() itself!

What I'd like to to is to avoid having to create my own mock object here and use some smarter approach here using unittest.mock

Upvotes: 0

Views: 649

Answers (1)

MrBean Bremen
MrBean Bremen

Reputation: 16815

What you are testing is basically what arguments run_update is called with. You can just mock the connection and use assert_called_xxx methods on the mock, and if you want to check specific arguments instead of all arguments you can check call_args on the mock. Here is an example that matches your sample code:

@mock.patch("some_module.MySql")
def test_proceed_to(mocked):
    # we need the mocked instance, not the class
    sql_mock = mocked.return_value  
    conn = sql_mock.connect()  # conn is now a mock - you also could have created a mock manually
    proceed_to(conn, 1, '2', 3, 4, 5)
    # assert that the function was called with the correct arguments
    conn.run_update.assert_called_once()
    # conn.run_update.assert_called_once_with() would compare all arguments
    assert conn.run_update.call_args[0][1] == [1, '2', 3, 4, 5]
    # call_args[0] are the positional arguments, so this checks the second positional argument

Upvotes: 1

Related Questions