def test_set_index_timezone():
    s_naive = pd.Series(pd.date_range('20130101', periods=3))
    s_aware = pd.Series(pd.date_range('20130101', periods=3, tz='US/Eastern'))
    df = pd.DataFrame({'tz': s_aware, 'notz': s_naive})
    d = dd.from_pandas(df, 2)
    d1 = d.set_index('notz', npartitions=1)
    s1 = pd.DatetimeIndex(s_naive.values, dtype=s_naive.dtype)
    assert (d1.divisions[0] == s_naive[0] == s1[0])
    assert (d1.divisions[(- 1)] == s_naive[2] == s1[2])
    d2 = d.set_index('tz', npartitions=1)
    s2 = pd.DatetimeIndex(s_aware, dtype=s_aware.dtype)
    assert (d2.divisions[0] == s2[0])
    assert (d2.divisions[(- 1)] == s2[2])
    assert (d2.divisions[0].tz == s2[0].tz)
    assert (d2.divisions[0].tz is not None)
    s2badtype = pd.DatetimeIndex(s_aware.values, dtype=s_naive.dtype)
    if PANDAS_GT_120:
        assert (not (d2.divisions[0] == s2badtype[0]))
    else:
        with pytest.raises(TypeError):
            (d2.divisions[0] == s2badtype[0])
def test_set_index_timezone():
    s_naive = pd.Series(pd.date_range('20130101', periods=3))
    s_aware = pd.Series(pd.date_range('20130101', periods=3, tz='US/Eastern'))
    df = pd.DataFrame({'tz': s_aware, 'notz': s_naive})
    d = dd.from_pandas(df, 2)
    d1 = d.set_index('notz', npartitions=1)
    s1 = pd.DatetimeIndex(s_naive.values, dtype=s_naive.dtype)
    assert (d1.divisions[0] == s_naive[0] == s1[0])
    assert (d1.divisions[(- 1)] == s_naive[2] == s1[2])
    d2 = d.set_index('tz', npartitions=1)
    s2 = pd.DatetimeIndex(s_aware, dtype=s_aware.dtype)
    assert (d2.divisions[0] == s2[0])
    assert (d2.divisions[(- 1)] == s2[2])
    assert (d2.divisions[0].tz == s2[0].tz)
    assert (d2.divisions[0].tz is not None)
    s2badtype = pd.DatetimeIndex(s_aware.values, dtype=s_naive.dtype)
    if PANDAS_GT_120:
        assert (not (d2.divisions[0] == s2badtype[0]))
    else:
        with pytest.raises(TypeError):
            (d2.divisions[0] == s2badtype[0])