Repository: pandas-dev/pandas
File: pandas/tests/frame/methods/test_join.py (Python)

Code before:
        ):
            df1.join(df2, on="a")

    def test_frame_join_tzaware(self):
        tz = zoneinfo.ZoneInfo("US/Central")
        test1 = DataFrame(
            np.zeros((6, 3)),
            index=date_range("2012-11-15 00:00:00", periods=6, freq="100ms", tz=tz),
        )
        test2 = DataFrame(
            np.zeros((3, 3)),
            index=date_range("2012-11-15 00:00:00", periods=3, freq="250ms", tz=tz),
            columns=range(3, 6),
        )

        result = test1.join(test2, how="outer")
        expected = test1.index.union(test2.index)

        tm.assert_index_equal(result.index, expected)
        assert result.index.tz.key == "US/Central"

    def test_frame_join_categorical_index(self):
        # GH 61675
        cat_data = pd.Categorical(
            [15, 16, 17, 18],
            categories=pd.Series(list(range(3, 24)), dtype="Int64"),
            ordered=True,
        )
        values1 = "a b c d".split()
        values2 = "xyzzy foo bar ...".split()
        df1 = DataFrame({"hr": cat_data, "values1": values1}).set_index("hr")
        df2 = DataFrame({"hr": cat_data, "values2": values2}).set_index("hr")
        df1.columns = pd.CategoricalIndex([4], dtype=cat_data.dtype, name="other_hr")
        df2.columns = pd.CategoricalIndex([3], dtype=cat_data.dtype, name="other_hr")

        df_joined_1 = (
            df1.reset_index(level="hr")
            .merge(df2.reset_index(level="hr"), on="hr")
            .set_index("hr")
        )
        expected1 = DataFrame(
            {"hr": cat_data, "values1": values1, "values2": values2}
        ).set_index("hr")
        expected1.columns = Index([4, 3], dtype="object", name="other_hr")

        tm.assert_frame_equal(df_joined_1, expected1)

        df_joined_2 = df1.join(df2)
        expected2 = DataFrame(
            {"hr": cat_data, "values1": values1, "values2": values2}
        ).set_index("hr")

Review comment (bug):
Can you reduce the size here; even 20 elements can make stepping through with a debugger harder. Shoot for 3 or 5 if necessary.

Code after:
        ):
            df1.join(df2, on="a")

    def test_frame_join_tzaware(self):
        tz = zoneinfo.ZoneInfo("US/Central")
        test1 = DataFrame(
            np.zeros((6, 3)),
            index=date_range("2012-11-15 00:00:00", periods=6, freq="100ms", tz=tz),
        )
        test2 = DataFrame(
            np.zeros((3, 3)),
            index=date_range("2012-11-15 00:00:00", periods=3, freq="250ms", tz=tz),
            columns=range(3, 6),
        )

        result = test1.join(test2, how="outer")
        expected = test1.index.union(test2.index)

        tm.assert_index_equal(result.index, expected)
        assert result.index.tz.key == "US/Central"

    def test_frame_join_categorical_index(self):
        # GH 61675
        cat_data = pd.Categorical(
            [3, 4],
            categories=pd.Series([2, 3, 4, 5], dtype="Int64"),
            ordered=True,
        )
        values1 = "a b".split()
        values2 = "foo bar".split()
        df1 = DataFrame({"hr": cat_data, "values1": values1}).set_index("hr")
        df2 = DataFrame({"hr": cat_data, "values2": values2}).set_index("hr")
        df1.columns = pd.CategoricalIndex([4], dtype=cat_data.dtype, name="other_hr")
        df2.columns = pd.CategoricalIndex([3], dtype=cat_data.dtype, name="other_hr")

        df_joined = df1.join(df2)
        expected = DataFrame(
            {"hr": cat_data, "values1": values1, "values2": values2}
        ).set_index("hr")
        expected.columns = pd.CategoricalIndex(
            [4, 3], dtype=cat_data.dtype, name="other_hr"
        )

        tm.assert_frame_equal(df_joined, expected)