from caffe2.python import core, workspace, test_util
import os
import shutil
import tempfile
import unittest


class CheckpointTest(test_util.TestCase):
    """A simple test case to make sure that the checkpoint behavior is correct.
    """

    @unittest.skipIf("LevelDB" not in core.C.registered_dbs(), "Need LevelDB")
    def testCheckpoint(self):
        temp_root = tempfile.mkdtemp()
        net = core.Net("test_checkpoint")
        # Note(jiayq): I am being a bit lazy here and am using the old iter
        # convention that does not have an input. Optionally change it to the
        # new style if needed.
        net.Iter([], "iter")
        net.ConstantFill([], "value", shape=[1, 2, 3])
        net.Checkpoint(["iter", "value"], [],
                       db=os.path.join(temp_root, "test_checkpoint_at_%05d"),
                       db_type="leveldb", every=10, absolute_path=True)
        self.assertTrue(workspace.CreateNet(net))
        for i in range(100):
            self.assertTrue(workspace.RunNet("test_checkpoint"))
        for i in range(1, 10):
            # Print statements are only for debugging purposes.
            # print("Asserting %d" % i)
            # print(os.path.join(temp_root, "test_checkpoint_at_%05d" % (i * 10)))
            self.assertTrue(os.path.exists(
                os.path.join(temp_root, "test_checkpoint_at_%05d" % (i * 10))))

        # Finally, clean up.
        shutil.rmtree(temp_root)


if __name__ == "__main__":
    unittest.main()