Method test_cached_reader

caffe2/python/dataio_test.py:383–420

(self)

    @unittest.skipIf("LevelDB" not in core.C.registered_dbs(), "Need LevelDB")
    def test_cached_reader(self):
        ws = workspace.C.Workspace()
        session = LocalSession(ws)
        db_path = self._make_temp_path()

        # Read data for the first time.
        cached_reader1 = CachedReader(
            self._build_source_reader(ws, 100), db_path, loop_over=False,
        )
        build_cache_step = cached_reader1.build_cache_step()
        session.run(build_cache_step)

        data = self._read_all_data(ws, cached_reader1, session)
        self.assertEqual(sorted(data), list(range(100)))

        # Read data from cache.
        cached_reader2 = CachedReader(
            self._build_source_reader(ws, 200), db_path,
        )
        build_cache_step = cached_reader2.build_cache_step()
        session.run(build_cache_step)

        data = self._read_all_data(ws, cached_reader2, session)
        self.assertEqual(sorted(data), list(range(100)))

        self._delete_path(db_path)

        # We removed cache so we expect to receive data from original reader.
        cached_reader3 = CachedReader(
            self._build_source_reader(ws, 300), db_path,
        )
        build_cache_step = cached_reader3.build_cache_step()
        session.run(build_cache_step)

        data = self._read_all_data(ws, cached_reader3, session)
        self.assertEqual(sorted(data), list(range(300)))

        self._delete_path(db_path)

    @unittest.skipIf("LevelDB" not in core.C.registered_dbs(), "Need LevelDB")
    def test_db_file_reader(self):