REF: Assert json roundtrip equal (#28626)

Author: bganglia

pandas/tests/io/json/test_pandas.py

@@ -37,6 +37,14 @@
 _mixed_frame = _frame.copy()
 
 
+def assert_json_roundtrip_equal(result, expected, orient):
+    if orient == "records" or orient == "values":
+        expected = expected.reset_index(drop=True)
+    if orient == "values":
+        expected.columns = range(len(expected.columns))
+    assert_frame_equal(result, expected)
+
+
 class TestPandasContainer:
     @pytest.fixture(scope="function", autouse=True)
     def setup(self, datapath):
@@ -90,25 +98,15 @@ def test_frame_double_encoded_labels(self, orient):
         result = read_json(df.to_json(orient=orient), orient=orient)
         expected = df.copy()
 
-        if orient == "records" or orient == "values":
-            expected = expected.reset_index(drop=True)
-        if orient == "values":
-            expected.columns = range(len(expected.columns))
-
-        assert_frame_equal(result, expected)
+        assert_json_roundtrip_equal(result, expected, orient)
 
     @pytest.mark.parametrize("orient", ["split", "records", "values"])
     def test_frame_non_unique_index(self, orient):
         df = DataFrame([["a", "b"], ["c", "d"]], index=[1, 1], columns=["x", "y"])
         result = read_json(df.to_json(orient=orient), orient=orient)
         expected = df.copy()
 
-        if orient == "records" or orient == "values":
-            expected = expected.reset_index(drop=True)
-        if orient == "values":
-            expected.columns = range(len(expected.columns))
-
-        assert_frame_equal(result, expected)
+        assert_json_roundtrip_equal(result, expected, orient)
 
     @pytest.mark.parametrize("orient", ["index", "columns"])
     def test_frame_non_unique_index_raises(self, orient):
@@ -172,12 +170,7 @@ def test_roundtrip_simple(self, orient, convert_axes, numpy, dtype):
             # TODO: debug why sort is required
             expected = expected.sort_index()
 
-        if orient == "records" or orient == "values":
-            expected = expected.reset_index(drop=True)
-        if orient == "values":
-            expected.columns = range(len(expected.columns))
-
-        tm.assert_frame_equal(result, expected)
+        assert_json_roundtrip_equal(result, expected, orient)
 
     @pytest.mark.parametrize("dtype", [False, np.int64])
     @pytest.mark.parametrize("convert_axes", [True, False])
@@ -191,11 +184,6 @@ def test_roundtrip_intframe(self, orient, convert_axes, numpy, dtype):
         if not numpy and (orient == "index" or (PY35 and orient == "columns")):
             expected = expected.sort_index()
 
-        if orient == "records" or orient == "values":
-            expected = expected.reset_index(drop=True)
-        if orient == "values":
-            expected.columns = range(len(expected.columns))
-
         if (
             numpy
             and (is_platform_32bit() or is_platform_windows())
@@ -205,7 +193,7 @@ def test_roundtrip_intframe(self, orient, convert_axes, numpy, dtype):
             # TODO: see what is causing roundtrip dtype loss
             expected = expected.astype(np.int32)
 
-        tm.assert_frame_equal(result, expected)
+        assert_json_roundtrip_equal(result, expected, orient)
 
     @pytest.mark.parametrize("dtype", [None, np.float64, np.int, "U3"])
     @pytest.mark.parametrize("convert_axes", [True, False])
@@ -246,12 +234,7 @@ def test_roundtrip_str_axes(self, orient, convert_axes, numpy, dtype):
         elif orient == "records" and convert_axes:
             expected.columns = expected.columns.astype(np.int64)
 
-        if orient == "records" or orient == "values":
-            expected = expected.reset_index(drop=True)
-        if orient == "values":
-            expected.columns = range(len(expected.columns))
-
-        tm.assert_frame_equal(result, expected)
+        assert_json_roundtrip_equal(result, expected, orient)
 
     @pytest.mark.parametrize("convert_axes", [True, False])
     @pytest.mark.parametrize("numpy", [True, False])
@@ -277,12 +260,7 @@ def test_roundtrip_categorical(self, orient, convert_axes, numpy):
         if not numpy and (orient == "index" or (PY35 and orient == "columns")):
             expected = expected.sort_index()
 
-        if orient == "records" or orient == "values":
-            expected = expected.reset_index(drop=True)
-        if orient == "values":
-            expected.columns = range(len(expected.columns))
-
-        tm.assert_frame_equal(result, expected)
+        assert_json_roundtrip_equal(result, expected, orient)
 
     @pytest.mark.parametrize("convert_axes", [True, False])
     @pytest.mark.parametrize("numpy", [True, False])
@@ -320,12 +298,7 @@ def test_roundtrip_timestamp(self, orient, convert_axes, numpy):
 
             expected.index = idx
 
-        if orient == "records" or orient == "values":
-            expected = expected.reset_index(drop=True)
-        if orient == "values":
-            expected.columns = range(len(expected.columns))
-
-        tm.assert_frame_equal(result, expected)
+        assert_json_roundtrip_equal(result, expected, orient)
 
     @pytest.mark.parametrize("convert_axes", [True, False])
     @pytest.mark.parametrize("numpy", [True, False])
@@ -354,12 +327,7 @@ def test_roundtrip_mixed(self, orient, convert_axes, numpy):
         if not numpy and (orient == "index" or (PY35 and orient == "columns")):
             expected = expected.sort_index()
 
-        if orient == "records" or orient == "values":
-            expected = expected.reset_index(drop=True)
-        if orient == "values":
-            expected.columns = range(len(expected.columns))
-
-        tm.assert_frame_equal(result, expected)
+        assert_json_roundtrip_equal(result, expected, orient)
 
     @pytest.mark.parametrize(
         "data,msg,orient",