fix typing in tests

epoch8 · Aug 31, 2024 · 3952083 · 3952083
1 parent 87224d2
commit 3952083
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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)