Jake
Jake

Reputation: 2912

Flask unit-test: Unable to mock a function in a post request

I have the following flask app.

# app.py

from flask import Flask, request
from predict import query_sku

app = Flask(__name__)

@app.route("/predict", methods=["POST"])
def predict():
    content = request.json
    max_results = content["resultSize"]
    input_sku_list = content["sku"]

    skus = query_sku(input_sku_list, max_results)
    return {"sku": skus}

if __name__ == "__main__":
    app.run()

I wrote a unit-test for it using pytest and tried to mock the query_sku function using unittest.mock.

import sys
from unittest.mock import Mock

import pytest

import app

def test_api_mocked_model():
    sys.modules["predict"] = Mock()
    from predict import query_sku
    query_sku.return_value = "dummy"

    with app.app.test_client() as client:
        response = client.post('/predict', json={"resultSize":10,"sku": "x"}).json

    assert response == {"sku": "dummy"}
    del sys.modules['predict']

But I was unable to mock that function within the request. It just gave the following assertion error.

>       assert response == {"sku": "dummy"}
E       AssertionError: assert None == {'sku': 'dummy'}
E         +None
E         -{'sku': 'dummy'}

tests/unit_tests/test_api.py:34: AssertionError

How can I get it to work?

[EDIT]

I added in the query_sku function below. Intentionally return a value that is different from the Mock function return_value.

# predict.py
def query_sku(input_sku_list, topn):
    return "actual function"

But the unit-test is still querying from the actual function, as shown below.

assert response == {"sku": "dummy"}
E       AssertionError: assert {'sku': 'actual function'} == {'sku': 'dummy'}
E         Differing items:
E         {'sku': 'actual function'} != {'sku': 'dummy'}

Upvotes: 0

Views: 2363

Answers (2)

Daniel Viglione
Daniel Viglione

Reputation: 9407

Actually, the answer above did not work for me. Rather than using MagicMock, I used @mock.patch and that worked. And it is less complicated too. For example, in your app.py, you can have an API endpoint that you want to stub out with a mock:

app.py

def fetch_nodes(usr, passwd, hostname, node_name):
  sftp_connection = get_sftp_connection(usr, passwd, hostname)
  nodes = sftp_connection.listdir_attr(node_name)
  return nodes

Now in test_app.py, you can stub it out like this:

test_app.py

from app import app
import unittest
from unittest import mock

class FlaskTest(unittest.TestCase):
  @mock.patch('app.fetch_nodes')
  def test_list_child_nodes(self, mocked):
    tester = app.test_client(self)

    mocked.return_value = []

    response = tester.post("/api/listchildnodes", json={
      'usr': 'test',
      'jwt': 'test',
      'node_name': 'test',
    })

    statuscode = response.status_code
    self.assertEqual(statuscode, 200)
  
if __name__ == "__main__":
  unittest.main()

The important thing is the definition of mock.patch('app.fetch_nodes') which references the fetch_nodes function in app.py and secondly it is passed to the test function as "mocked". Then we set the return value of "mocked". And finally call the flask endpoint. Now the flask endpoint will use the mock instead of actually hitting an sftp server.

Upvotes: 0

Bao Tran
Bao Tran

Reputation: 626

from unittest.mock import MagicMock

def test_api_mocked_model():
     ## sys.modules["predict"] = Mock() ## why sys module ?
    from predict import query_sku
    query_sku = MagicMock(return_value="dummy") # mock directly 
  


    with app.app.test_client() as client:
        response = client.post('/predict', json={"resultSize":10,"sku": "x"}).json

    assert response == {"sku": "dummy"}
    del sys.modules['predict']

Could you try this code ?

Upvotes: 1

Related Questions