test: fix test

Author: mamo3gr

gokart/target.py

@@ -80,11 +80,12 @@ def __init__(
         target: luigi.target.FileSystemTarget,
         processor: FileProcessor,
         task_lock_params: TaskLockParams,
-        expected_dataframe_type: Optional[pa.DataFrameModel] = None,
+        expected_dataframe_type: Optional[type[pa.DataFrameModel]] = None,
     ) -> None:
         self._target = target
         self._processor = processor
         self._task_lock_params = task_lock_params
+        self._expected_dataframe_type = expected_dataframe_type
 
     def _exists(self) -> bool:
         return self._target.exists()
@@ -95,14 +96,14 @@ def _get_task_lock_params(self) -> TaskLockParams:
     def _load(self) -> Any:
         with self._target.open('r') as f:
             obj = self._processor.load(f)
-            if self.expected_dataframe_type is not None:
-                return self.expected_dataframe_type(obj)
+            if self._expected_dataframe_type is not None:
+                return pa.typing.DataFrame[self._expected_dataframe_type](obj)
 
             return obj
 
     def _dump(self, obj) -> None:
-        if self.expected_dataframe_type is not None:
-            self.expected_dataframe_type.validate(obj)
+        if self._expected_dataframe_type is not None:
+            self._expected_dataframe_type.validate(obj)
 
         with self._target.open('w') as f:
             self._processor.dump(obj, f)

test/test_target.py

@@ -14,6 +14,7 @@
 from matplotlib import pyplot
 from moto import mock_s3
 
+from gokart.conflict_prevention_lock.task_lock import make_task_lock_params
 from gokart.file_processor import _ChunkedLargeFileReader, make_file_processor
 from gokart.target import SingleFileTarget, make_model_target, make_target
 
@@ -263,22 +264,26 @@ def test_typed_target(self):
         test_case = pd.DataFrame(dict(a=[1, 2]))
 
         with tempfile.TemporaryDirectory() as temp_dir:
-            _task_lock_params = None
             file_path = os.path.join(temp_dir, 'test.csv')
+            unique_id = 'test_unique_id'
+            _task_lock_params = make_task_lock_params(file_path=file_path, unique_id=unique_id)
             processor = make_file_processor(file_path, store_index_in_feather=False)
             file_system_target = luigi.LocalTarget(file_path, format=processor.format())
             file_target = SingleFileTarget(target=file_system_target, processor=processor, task_lock_params=_task_lock_params, expected_dataframe_type=self.DummyDataFrameSchema)
 
             file_target.dump(test_case)
             dumped_data = file_target.load()
-            self.assertIsInstance(dumped_data, self.DummyDataFrameSchema)
+            self.assertIsInstance(dumped_data, pa.typing.DataFrame)
+            self.DummyDataFrameSchema.validate(dumped_data)
+
 
     def test_invalid_typed_target(self):
         test_case = pd.DataFrame(dict(a=['1', '2']))
 
         with tempfile.TemporaryDirectory() as temp_dir:
-            _task_lock_params = None
             file_path = os.path.join(temp_dir, 'test.csv')
+            unique_id = 'test_unique_id'
+            _task_lock_params = make_task_lock_params(file_path=file_path, unique_id=unique_id)
             processor = make_file_processor(file_path, store_index_in_feather=False)
             file_system_target = luigi.LocalTarget(file_path, format=processor.format())
             file_target = SingleFileTarget(target=file_system_target, processor=processor, task_lock_params=_task_lock_params, expected_dataframe_type=self.DummyDataFrameSchema)