How do I mock a function which has been called twice in the same file, with different parameters? Here is my code snippet:

code_file.py

from prefect import Client
client = Client(
    api_server = <ip of server>
)
def run_flow(name1, name2):
     # some graphql queries, not shown here
     value1 = client.graphql(graphql_value1, variables={'name': name1})
     print("First client call return value:- ", value1)
     value2 = client.graphql(graphql_value2, variables={'name': name2, 'id': value1['data'][0]['id']})
     print("Second client call return value:- ", value2)

     run_id = client.create_flow(id = value2['data'][0]['id'])
        return run_id

code_test.py

# I want to mock these client calls.
# However when I am mocking, only the first client.graphql function is being mocked. 
# How do I mock the other client.graphql and the client.create_flow function?

import unittest
import pytest
from pytest_mock import mocker
from unittest.mock import Mock
from fastapi.testclient import TestClient
from app.main import app

client = TestClient(app)

class test_flow(unittest.TestCase):
 @patch("code_file.client.graphql", side_effect=["test1s", "test2s"])
 @patch("code_file.client.create_flow", return_value="test_id")
 def test_flow(self, s1):
   # Below post call invokes run_flow function
   response = client.post("/name1/name2")
   assert run_flow("name1", "name2") == "test_id"

First mock call to graphql is succeeding. The second graphql call is not getting mocked. It is trying to contact the actual server and receiving 404. How can I mock both graphql client calls?

CodePudding user response：

I tried to reproduce your codes, but for me client.graphql got mocked perfectly.

Here are my codes.

Folder structure

├── code_file.py
├── code_test.py
└── prefect.py

# prefect.py
class Client:
    def __init__(self, *arg, **kwargs):
        ...

    def graphql(self, *args, **kwargs):
        print(f"graphql got called with args: {args}, kwargs: {kwargs}")

    def create_flow(self, *args, **kwargs):
        print(f"create_flow got called with args: {args}, kwargs: {kwargs}")

# code_file.py
from prefect import Client

graphql_value1 = "graphql_value1"
graphql_value2 = "graphql_value2"

client = Client(api_server="http://example.com")


def run_flow(name1, name2):
    # some graphql queries, not shown here
    value1 = client.graphql(graphql_value1, variables={"name": name1})
    print("First client call return value:- ", value1)
    value2 = client.graphql(graphql_value2, variables={"name": name2, "id": value1["data"][0]["id"]})
    print("Second client call return value:- ", value2)

    run_id = client.create_flow(id=value2["data"][0]["id"])
    return run_id

# code_test.py
import unittest
from unittest.mock import patch, MagicMock, call
from code_file import run_flow


class TestFlow(unittest.TestCase):
    @patch("code_file.client.graphql")
    @patch("code_file.client.create_flow")
    def test_flow(self, create_flow: MagicMock, graphql: MagicMock) -> None:
        first_graphql_return_value = {"data": [{"id": 1}]}
        second_graphl_return_value = {"data": [{"id": 2}]}
        graphql.side_effect = [
            first_graphql_return_value,
            second_graphl_return_value,
        ]
        create_flow.return_value = "test_id"

        self.assertEqual(run_flow("name1", "name2"), "test_id")

        create_flow.assert_called_once_with(id=2)
        graphql.assert_has_calls(
            [
                call("graphql_value1", variables={"name": "name1"}),
                call("graphql_value2", variables={"name": "name2", "id": 1})
            ]
        )

Running unittest by using the command

python -m unittest code_test.py

produces the following output

First client call return value:-  {'data': [{'id': 1}]}
Second client call return value:-  {'data': [{'id': 2}]}
.
----------------------------------------------------------------------
Ran 1 test in 0.001s

You'll see that the prints in prefect.Client methods do not get printed.

Original answer

Explanation

From python official documentation

If side_effect is an iterable then each call to the mock will return the next value from the iterable.

When you set side_effect argument with an iterable, the mock will return next value from the iterable. More detail can be found here

Solution

def test_flow(mocker):
    # Some lines of code, not shown here
    m = mocker.patch("code_file.client.graphql", side_effect=["value1", "value2"])
    ...