def test_wf_typed_schema():
    schema1 = FlyteSchema[kwtypes(x=int, y=str)]

    @task
    def t1() -> schema1:
        s = schema1()
        s.open().write(pandas.DataFrame(data={'x': [1, 2], 'y': ['3', '4']}))
        return s

    @task
    def t2(s: FlyteSchema[kwtypes(x=int, y=str)]) -> FlyteSchema[kwtypes(x=int)]:
        df = s.open().all()
        return df[s.column_names()[:(- 1)]]

    @workflow
    def wf() -> FlyteSchema[kwtypes(x=int)]:
        return t2(s=t1())
    w = t1()
    assert (w is not None)
    df = w.open(override_mode=SchemaOpenMode.READ).all()
    result_df = (df.reset_index(drop=True) == pandas.DataFrame(data={'x': [1, 2], 'y': ['3', '4']}).reset_index(drop=True))
    assert result_df.all().all()
    df = t2(s=w.as_readonly())
    assert (df is not None)
    result_df = (df.reset_index(drop=True) == pandas.DataFrame(data={'x': [1, 2]}).reset_index(drop=True))
    assert result_df.all().all()
    x = wf()
    df = x.open().all()
    result_df = (df.reset_index(drop=True) == pandas.DataFrame(data={'x': [1, 2]}).reset_index(drop=True))
    assert result_df.all().all()
def test_wf_typed_schema():
    schema1 = FlyteSchema[kwtypes(x=int, y=str)]

    @task
    def t1() -> schema1:
        s = schema1()
        s.open().write(pandas.DataFrame(data={'x': [1, 2], 'y': ['3', '4']}))
        return s

    @task
    def t2(s: FlyteSchema[kwtypes(x=int, y=str)]) -> FlyteSchema[kwtypes(x=int)]:
        df = s.open().all()
        return df[s.column_names()[:(- 1)]]

    @workflow
    def wf() -> FlyteSchema[kwtypes(x=int)]:
        return t2(s=t1())
    w = t1()
    assert (w is not None)
    df = w.open(override_mode=SchemaOpenMode.READ).all()
    result_df = (df.reset_index(drop=True) == pandas.DataFrame(data={'x': [1, 2], 'y': ['3', '4']}).reset_index(drop=True))
    assert result_df.all().all()
    df = t2(s=w.as_readonly())
    assert (df is not None)
    result_df = (df.reset_index(drop=True) == pandas.DataFrame(data={'x': [1, 2]}).reset_index(drop=True))
    assert result_df.all().all()
    x = wf()
    df = x.open().all()
    result_df = (df.reset_index(drop=True) == pandas.DataFrame(data={'x': [1, 2]}).reset_index(drop=True))
    assert result_df.all().all()