I'm trying to test a method which retrieves a mysql connection inside a with block but when i try to get a mocked cursor from it and also the result of the execute from the cursor is not properly mocked.
Here is the method to get the connection
def has_pending_task(self, host, port):
with self.get_mysql_connection(host, port) as conn:
cursor = conn.cursor()
table = self.config.get_table()
cursor.execute(
'select count(*) from {} where status in ("PENDING", "RUNNING") and host=%s'.format(
table), host)
count = cursor.fetchone()
cursor.close()
return count > 0
@contextmanager
def get_mysql_connection(self, host, port):
db_credentials = self.config.get_credentials()
schema = self.config.get_schema()
if not db_credentials:
logging.error(
"There was an error getting db info for cluster {}. Please review the config".format(self.cluster))
return None
conn = None
try:
conn = MySQLConnection(user=db_credentials[0], password=db_credentials[1],
host=host, port=port, database=schema)
except Error as error:
logging.error("Couldn't create a mysql connection: {}".format(error.msg))
return conn
the test looks like this:
def test_has_no_pending_tasks(self):
with mock.patch("my_package.my_class.Config") as config_mock:
config = config_mock.return_value
config_mock.get_instance.return_value = config
config.get_table.return_value = "table"
with mock.patch("mysql.connector.connection.MySQLConnection") as mysql_mock:
mysql = mysql_mock.return_value
with mock.patch('package.MyClass.get_mysql_connection') as conn_mock:
conn_mock.return_value = mysql
cursor = mysql_mock.cursor.return_value
cursor.fetchone.return_value = 0
Orchestrator("host", 3306).has_pending_tasks("host", 3306)
conn_mock.assert_called_once_with("host", 3306)
the problem is that as it's not mocking things properly, when does the return from the tested function, count is an instance of MagicMock and makes the test crash.
Any idea of what part i'm missing?
Thanks!
