From 69f6b843da1fe38a6b1bf1e6af698282fc491680 Mon Sep 17 00:00:00 2001
From: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Date: Sat, 23 Jul 2022 04:53:59 -0700
Subject: [PATCH] [RLlib] Test output length in DatasetReader with default
 IOContext. (#26852)

---
 rllib/offline/tests/test_dataset_reader.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/rllib/offline/tests/test_dataset_reader.py b/rllib/offline/tests/test_dataset_reader.py
index 77ec567df..9cb9809dc 100644
--- a/rllib/offline/tests/test_dataset_reader.py
+++ b/rllib/offline/tests/test_dataset_reader.py
@@ -139,7 +139,8 @@ class TestDatasetReader(unittest.TestCase):
             {"input": "dataset", "input_config": input_config}
         )
         reader = DatasetReader(dataset)
-        reader.next()
+        # Reads in one line of Pendulum dataset with 600 timesteps
+        assert len(reader.next()) == 600
 
 
 class TestUnzipIfNeeded(unittest.TestCase):