From 8ceafe556b635bfaced4d7a64e3a9c9c7862504e Mon Sep 17 00:00:00 2001
From: proxodilka <zeeron209@gmail.com>
Date: Wed, 25 May 2022 19:40:25 +0300
Subject: [PATCH] Rewrite test for issue#4456

Signed-off-by: proxodilka <zeeron209@gmail.com>
---
 modin/pandas/test/dataframe/test_indexing.py | 91 +++++++++++++-------
 1 file changed, 61 insertions(+), 30 deletions(-)

diff --git a/modin/pandas/test/dataframe/test_indexing.py b/modin/pandas/test/dataframe/test_indexing.py
index a99967eea94..c0854aae344 100644
--- a/modin/pandas/test/dataframe/test_indexing.py
+++ b/modin/pandas/test/dataframe/test_indexing.py
@@ -42,6 +42,7 @@
 from modin.config import NPartitions
 from modin.utils import get_current_execution
 from modin.test.test_utils import warns_that_defaulting_to_pandas
+from modin.pandas.indexing import is_range_like
 
 NPartitions.put(4)
 
@@ -402,43 +403,73 @@ def test_loc(data):
         modin_df.loc["NO_EXIST"]
 
 
-@pytest.mark.parametrize("equal_axes", [False, True])
-@pytest.mark.parametrize("axis", [0, 1])
-def test_loc_4456_1(axis, equal_axes):
-    if equal_axes and axis == 1:
-        # case equal to `equal_axes and axis == 0`
-        pass
-
+@pytest.mark.parametrize(
+    "key_getter, value_getter",
+    [
+        pytest.param(
+            lambda df, axis: (
+                (slice(None), df.axes[axis][:2])
+                if axis
+                else (df.axes[axis][:2], slice(None))
+            ),
+            lambda df, axis: df.iloc[:, :1] if axis else df.iloc[:1, :],
+            id="len(key)_>_len(value)",
+        ),
+        pytest.param(
+            lambda df, axis: (
+                (slice(None), df.axes[axis][:2])
+                if axis
+                else (df.axes[axis][:2], slice(None))
+            ),
+            lambda df, axis: df.iloc[:, :3] if axis else df.iloc[:3, :],
+            id="len(key)_<_len(value)",
+        ),
+        pytest.param(
+            lambda df, axis: (
+                (slice(None), df.axes[axis][:2])
+                if axis
+                else (df.axes[axis][:2], slice(None))
+            ),
+            lambda df, axis: df.iloc[:, :2] if axis else df.iloc[:2, :],
+            id="len(key)_==_len(value)",
+        ),
+    ],
+)
+@pytest.mark.parametrize("key_axis", [0, 1])
+@pytest.mark.parametrize("reverse_value_index", [True, False])
+@pytest.mark.parametrize("reverse_value_columns", [True, False])
+def test_loc_4456(
+    key_getter, value_getter, key_axis, reverse_value_index, reverse_value_columns
+):
     data = test_data["float_nan_data"]
     modin_df, pandas_df = pd.DataFrame(data), pandas.DataFrame(data)
 
-    # prepare `loc` operands
-    key = pandas_df.columns[:2]
-    index, columns = pandas_df.index, key
-    value = np.random.randint(0, 100, size=(pandas_df.shape[0], len(columns)))
-    if not equal_axes:
-        if axis == 0:
-            index = list(reversed(index))
-        else:
-            columns = list(reversed(columns))
+    key = key_getter(pandas_df, key_axis)
 
-    loc_key = (slice(None), key)
+    # `df.loc` doesn't work right for range-like indexers. Converting them to a list.
+    # https://github.com/modin-project/modin/issues/4497
+    if is_range_like(key[0]):
+        key = (list(key[0]), key[1])
+    if is_range_like(key[1]):
+        key = (key[0], list(key[1]))
 
-    # pandas DataFrame
-    df_value = pandas.DataFrame(value, index=index, columns=columns)
-    eval_loc(modin_df, pandas_df, df_value, loc_key)
+    value = pandas.DataFrame(
+        np.random.randint(0, 100, size=pandas_df.shape),
+        index=pandas_df.index,
+        columns=pandas_df.columns,
+    )
+    pdf_value = value_getter(value, key_axis)
+    mdf_value = value_getter(pd.DataFrame(value), key_axis)
 
-    # modin DataFrame
-    mdf_value = pd.DataFrame(df_value)
-    eval_loc(modin_df, pandas_df, (mdf_value, df_value), loc_key)
+    if reverse_value_index:
+        pdf_value = pdf_value.reindex(index=pdf_value.index[::-1])
+        mdf_value = mdf_value.reindex(index=mdf_value.index[::-1])
+    if reverse_value_columns:
+        pdf_value = pdf_value.reindex(columns=pdf_value.columns[::-1])
+        mdf_value = mdf_value.reindex(columns=mdf_value.columns[::-1])
 
-    if axis == 0:
-        # pandas Series
-        pd_series_value = df_value[df_value.columns[0]]
-        eval_loc(modin_df, pandas_df, pd_series_value, loc_key)
-        # modin Series
-        md_series_value = mdf_value[mdf_value.columns[0]]
-        eval_loc(modin_df, pandas_df, (md_series_value, pd_series_value), loc_key)
+    eval_loc(modin_df, pandas_df, pdf_value, key)
+    eval_loc(modin_df, pandas_df, (mdf_value, pdf_value), key)
 
 
 # This tests the bug from https://github.com/modin-project/modin/issues/3736