TST: IntegerNA Support for DataFrame.diff() (#34889)

Author: Marvzinc

pandas/tests/frame/methods/test_diff.py

@@ -169,3 +169,48 @@ def test_diff_sparse(self):
         )
 
         tm.assert_frame_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "axis,expected",
+        [
+            (
+                0,
+                pd.DataFrame(
+                    {
+                        "a": [np.nan, 0, 1, 0, np.nan, np.nan, np.nan, 0],
+                        "b": [np.nan, 1, np.nan, np.nan, -2, 1, np.nan, np.nan],
+                        "c": np.repeat(np.nan, 8),
+                        "d": [np.nan, 3, 5, 7, 9, 11, 13, 15],
+                    },
+                    dtype="Int64",
+                ),
+            ),
+            (
+                1,
+                pd.DataFrame(
+                    {
+                        "a": np.repeat(np.nan, 8),
+                        "b": [0, 1, np.nan, 1, np.nan, np.nan, np.nan, 0],
+                        "c": np.repeat(np.nan, 8),
+                        "d": np.repeat(np.nan, 8),
+                    },
+                    dtype="Int64",
+                ),
+            ),
+        ],
+    )
+    def test_diff_integer_na(self, axis, expected):
+        # GH#24171 IntegerNA Support for DataFrame.diff()
+        df = pd.DataFrame(
+            {
+                "a": np.repeat([0, 1, np.nan, 2], 2),
+                "b": np.tile([0, 1, np.nan, 2], 2),
+                "c": np.repeat(np.nan, 8),
+                "d": np.arange(1, 9) ** 2,
+            },
+            dtype="Int64",
+        )
+
+        # Test case for default behaviour of diff
+        result = df.diff(axis=axis)
+        tm.assert_frame_equal(result, expected)