pythonpandasdataframepython-unittestpython-unittest.mock

How to patch function within pandas DataFrame.apply call


I'm trying to write a test for a class which performs a pandas apply. Here's a simplified version:

import pandas as pd
class Foo:
    def bar(self, row):
        return "BAR"

    def apply(self, df: pd.DataFrame):
        df["value"] = df.apply(self.bar, axis=1)

Here's the test I've attempted to write:

import unittest
from unittest.mock import patch

from my_module import Foo

class TestFoo(unittest.TestCase):
    def test_apply_1(self):  # passes
        foo = Foo()
        df = pd.DataFrame([1])
        foo.apply(df)
        self.assertEqual(df["value"].values[0], "BAR")

    def test_apply_2(self):  # fails
        foo = Foo()
        df = pd.DataFrame([1])  
        with patch.object(Foo, "bar") as mock_bar:              
            mock_bar.return_value = "BAZ"
            foo.apply(df)
            self.assertEqual(df["value"].values[0], "BAZ")

    def test_apply_3(self):  # fails
        foo = Foo()
        df = pd.DataFrame([1])
        with patch.object(foo, "bar") as mock_bar:
            mock_bar.return_value = "BAZ"
            foo.apply(df)
            self.assertEqual(df["value"].values[0], "BAZ")

When I run this test, test_apply_2 and test_apply_3 fail with error:

ValueError: No objects to concatenate

which indicates to me that the patch isn't working properly. What am I getting wrong here?


Solution

  • I figured it out; it's due to https://github.com/pandas-dev/pandas/issues/45298 ; the right way to test would be:

    import unittest
    from unittest.mock import patch, Mock
    
    from my_module import Foo
    
    class TestFoo(unittest.TestCase):
        @patch.object(Foo, "bar", return_value="BAZ", new_callable=Mock)
        def test_apply_2(self):  # succeeds
            foo = Foo()
            df = pd.DataFrame([1])  
            foo.apply(df)
            self.assertEqual(df["value"].values[0], "BAZ")