diff --git a/test/unit_tests/readers/test_decord_reader.py b/test/unit_tests/readers/test_decord_reader.py index ab81290fc..f6fd7138c 100644 --- a/test/unit_tests/readers/test_decord_reader.py +++ b/test/unit_tests/readers/test_decord_reader.py @@ -72,7 +72,7 @@ def test_should_sample_only_iframe(self): batches = list(video_loader.read()) expected = self._batches_to_reader_convertor( - create_dummy_batches(filters=[i for i in range(0, NUM_FRAMES, k)]) + create_dummy_batches(filters=[i for i in range(0, NUM_FRAMES, k)], is_from_storage=True) ) print(batches[0].frames) @@ -96,7 +96,7 @@ def test_should_sample_every_k_frame_with_predicate(self): value = NUM_FRAMES // 2 start = value + k - (value % k) if value % k else value expected = self._batches_to_reader_convertor( - create_dummy_batches(filters=[i for i in range(start, NUM_FRAMES, k)]) + create_dummy_batches(filters=[i for i in range(start, NUM_FRAMES, k)], is_from_storage=True) ) self.assertEqual(batches, expected) @@ -123,7 +123,7 @@ def test_should_sample_every_k_frame_with_predicate(self): batches = list(video_loader.read()) start = value + k - (value % k) if value % k else value expected = self._batches_to_reader_convertor( - create_dummy_batches(filters=[i for i in range(start, 8, k)]) + create_dummy_batches(filters=[i for i in range(start, 8, k)], is_from_storage=True) ) self.assertEqual(batches, expected) @@ -132,7 +132,7 @@ def test_should_return_one_batch(self): file_url=self.video_file_url, ) batches = list(video_loader.read()) - expected = self._batches_to_reader_convertor(create_dummy_batches()) + expected = self._batches_to_reader_convertor(create_dummy_batches(is_from_storage=True)) self.assertEqual(batches, expected) def test_should_return_batches_equivalent_to_number_of_frames(self): @@ -141,7 +141,7 @@ def test_should_return_batches_equivalent_to_number_of_frames(self): batch_mem_size=self.frame_size, ) batches = list(video_loader.read()) - expected = self._batches_to_reader_convertor(create_dummy_batches(batch_size=1)) + expected = self._batches_to_reader_convertor(create_dummy_batches(batch_size=1, is_from_storage=True)) self.assertEqual(batches, expected) def test_should_sample_every_k_frame(self): @@ -152,7 +152,7 @@ def test_should_sample_every_k_frame(self): ) batches = list(video_loader.read()) expected = self._batches_to_reader_convertor( - create_dummy_batches(filters=[i for i in range(0, NUM_FRAMES, k)]) + create_dummy_batches(filters=[i for i in range(0, NUM_FRAMES, k)], is_from_storage=True) ) self.assertEqual(batches, expected) diff --git a/test/util.py b/test/util.py index 4ffb2d9a2..4e759b4f0 100644 --- a/test/util.py +++ b/test/util.py @@ -451,6 +451,7 @@ def create_dummy_batches( batch_size=10, start_id=0, video_dir=None, + is_from_storage=False, # if cover test directly from storage, it needs to append a _row_number ): video_dir = video_dir or get_tmp_dir() @@ -461,7 +462,6 @@ def create_dummy_batches( data.append( { "myvideo._row_id": 1, - "myvideo._row_number": i + start_id, "myvideo.name": os.path.join(video_dir, "dummy.avi"), "myvideo.id": i + start_id, "myvideo.data": np.array( @@ -471,6 +471,9 @@ def create_dummy_batches( } ) + if is_from_storage: + data[-1]["myvideo._row_number"] = i + start_id + if len(data) % batch_size == 0: yield Batch(pd.DataFrame(data)) data = []