From f8dde26c13ddded14ecf9f497f0cd32971656cf6 Mon Sep 17 00:00:00 2001
From: Jiashen Cao <caojiashen24@gmail.com>
Date: Thu, 7 Sep 2023 22:21:05 -0400
Subject: [PATCH] update test

---
 test/unit_tests/readers/test_decord_reader.py | 12 ++++++------
 test/util.py                                  |  5 ++++-
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/test/unit_tests/readers/test_decord_reader.py b/test/unit_tests/readers/test_decord_reader.py
index ab81290fc1..f6fd7138c7 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 4ffb2d9a2c..4e759b4f02 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 = []