Mocking a function within a mapped function

Ray Libraries (Data, Train, Tune, Serve) Ray Data
1

How severe does this issue affect your experience of using Ray?

  • High: It blocks me to complete my task.

While unit-testing func I want to mock either func1 or func2, but I have been unable to do so.
If I test map_fn and mock the above functions, it works, though.

Here’s a toy example to reproduce the issue:

module.py

import ray.data as ray_data

Item = list[int]

def func1(item: Item) -> Item:
   assert False, "We want to mock this function because this will fail"
   return [i + 1 for i in item]

def func2(item: Item) -> Item:
   assert False, "We want to mock this function because this will fail"
   return [i + 3 for i in item]

def map_fn(item: Item) -> Item:
  print(func1)
  print(func2)

  out = func1(item)
  out = func2(out)

  return out


def func(dataset: ray_data.Dataset):
  return dataset.flat_map(map_fn)

test_module.py

import unittest
from unittest.mock import patch

import ray.data as ray_data

from module import map_fn, func


class TestFunc(unittest.TestCase):
    @patch('module.func2')
    @patch('module.func1')
    def test(self, mock_func1, mock_func2):
        mock_func1.return_value = [1, 2, 3, 4, 5]
        mock_func2.return_value = [1, 2, 3, 4, 5]

        item1 = [1, 2, 3]
        item2 = [4, 5, 6]

        dataset = ray_data.from_items([item1, item2])
        output = func(dataset)

        num = output.count()
        actual_output = output.take()


class TestMapFn(unittest.TestCase):
    @patch('module.func2')
    @patch('module.func1')
    def test(self, mock_func1, mock_func2):
        mock_func1.return_value = [1, 2, 3, 4, 5]
        mock_func2.return_value = [1, 2, 3, 4, 5]

        item1 = [1, 2, 3]

        output = map_fn(item1)

        self.assertListEqual(output, [1, 2, 3, 4, 5])


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

$ python test_module.py TestMapFn
<MagicMock name='func1' id='140208034068800'>
<MagicMock name='func2' id='140208034121760'>
.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

$ python test_module.py TestFunc
(_map_task pid=79123) <function func1 at 0x7fa10c555280>
(_map_task pid=79123) <function func2 at 0x7fa10c555310>
(_map_task pid=79121) <function func1 at 0x7f38fc5d6280>
(_map_task pid=79121) <function func2 at 0x7f38fc5d6310>
FAILED

How should I proceed here?