diff --git a/tests/test_complex_pipeline.py b/tests/test_complex_pipeline.py index b839438c..995d98d3 100644 --- a/tests/test_complex_pipeline.py +++ b/tests/test_complex_pipeline.py @@ -1,3 +1,5 @@ +from typing import cast + import pandas as pd import pytest from sqlalchemy import Column @@ -154,7 +156,7 @@ def complex_function( TEST__PIPELINE, TEST__PREDICTION, TEST__KEYPOINT, - idx=pd.DataFrame(columns=["item_id", "pipeline_id"]), + idx=cast(IndexDF, pd.DataFrame(columns=["item_id", "pipeline_id"])), ) run_steps(ds, steps) assert_datatable_equal(ds.get_table("output"), TEST_RESULT) diff --git a/tests/test_core_steps1.py b/tests/test_core_steps1.py index 56b6778f..e16b8831 100644 --- a/tests/test_core_steps1.py +++ b/tests/test_core_steps1.py @@ -83,7 +83,7 @@ def func(): return TEST_DF with pytest.raises(Exception): - do_batch_generate(func=func, ds=ds, output_dts=[tbl1]) + do_batch_generate(func=func, ds=ds, output_dts=[tbl1]) # type: ignore def test_inc_process_modify_values(dbconn) -> None: diff --git a/tests/test_core_steps2.py b/tests/test_core_steps2.py index 7a4be32e..26b54022 100644 --- a/tests/test_core_steps2.py +++ b/tests/test_core_steps2.py @@ -186,7 +186,7 @@ def test_batch_transform_with_dt_on_input_and_output(dbconn): "tbl2", table_store=TableStoreDB(dbconn, "tbl2_data", TEST_SCHEMA1, True) ) - df2 = TEST_DF1_1.loc[range(3, 8)] + df2 = TEST_DF1_1.loc[3:8] df2["a"] = df2["a"].apply(lambda x: x + 10) tbl1.store_chunk(TEST_DF1_1, now=0)