Method test_db_file_reader

caffe2/python/dataio_test.py:423–445

(self)

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

        # Build a cache DB file.
        cached_reader = CachedReader(
            self._build_source_reader(ws, 100),
            db_path=db_path,
            db_type='LevelDB',
        )
        build_cache_step = cached_reader.build_cache_step()
        session.run(build_cache_step)

        # Read data from cache DB file.
        db_file_reader = DBFileReader(
            db_path=db_path,
            db_type='LevelDB',
        )
        data = self._read_all_data(ws, db_file_reader, session)
        self.assertEqual(sorted(data), list(range(100)))

        self._delete_path(db_path)