# Path: fedjax/training/checkpoint.py
# _CHECKPOINT_PREFIX = 'checkpoint_'
# def _get_checkpoint_paths(base_path: str) -> List[str]:
#   def sort_key(path):
# def load_latest_checkpoint(root_dir: str) -> Optional[Tuple[Any, int]]:
# def save_checkpoint(root_dir: str,
#                     state: Any,
#                     round_num: int = 0,
#                     keep: int = 1):


import os
import jax.numpy as jnp
import numpy.testing as npt
from absl.testing import absltest
from fedjax.training import checkpoint
    for i in range(3):
      checkpoint.save_checkpoint(self._root_dir, state, round_num=i, keep=2)

    self.assertCountEqual(
        os.listdir(self._root_dir),
        ['checkpoint_00000001', 'checkpoint_00000002'])

  def test_get_checkpoint_paths(self):
    self._write_file('checkpoint_00000100')
    self._write_file('checkpoint_00000000')
    self._write_file('checkpoint_00000003')
    # Similarly named files that shouldn't be counted as valid checkpoints.
    self._write_file('checkpoint_haha')
    self._write_file('checkpoint_0ha')
    self._write_file('checkpoint_00000004.tf.ckpt')

    checkpoint_paths = checkpoint._get_checkpoint_paths(
        os.path.join(self._root_dir, 'checkpoint_'))

    basenames = [os.path.basename(cp) for cp in checkpoint_paths]
    self.assertEqual(
        basenames,
        ['checkpoint_00000000', 'checkpoint_00000003', 'checkpoint_00000100'])

  def test_load_latest_checkpoint(self):
    state_1 = {'w': jnp.ones(10)}
    state_2 = {'w': jnp.ones(10) * 2}
    checkpoint.save_checkpoint(self._root_dir, state_1, round_num=1, keep=2)
    checkpoint.save_checkpoint(self._root_dir, state_2, round_num=2, keep=2)


# Path: fedjax/training/checkpoint.py
# _CHECKPOINT_PREFIX = 'checkpoint_'
# def _get_checkpoint_paths(base_path: str) -> List[str]:
#   def sort_key(path):
# def load_latest_checkpoint(root_dir: str) -> Optional[Tuple[Any, int]]:
# def save_checkpoint(root_dir: str,
#                     state: Any,
#                     round_num: int = 0,
#                     keep: int = 1):


# Path: fedjax/training/checkpoint_test.py
import os
import jax.numpy as jnp
import numpy.testing as npt
from absl.testing import absltest
from fedjax.training import checkpoint

    for i in range(3):
      checkpoint.save_checkpoint(self._root_dir, state, round_num=i, keep=2)

    self.assertCountEqual(
        os.listdir(self._root_dir),
        ['checkpoint_00000001', 'checkpoint_00000002'])

  def test_get_checkpoint_paths(self):
    self._write_file('checkpoint_00000100')
    self._write_file('checkpoint_00000000')
    self._write_file('checkpoint_00000003')
    # Similarly named files that shouldn't be counted as valid checkpoints.
    self._write_file('checkpoint_haha')
    self._write_file('checkpoint_0ha')
    self._write_file('checkpoint_00000004.tf.ckpt')

    checkpoint_paths = checkpoint._get_checkpoint_paths(
        os.path.join(self._root_dir, 'checkpoint_'))

    basenames = [os.path.basename(cp) for cp in checkpoint_paths]
    self.assertEqual(
        basenames,
        ['checkpoint_00000000', 'checkpoint_00000003', 'checkpoint_00000100'])

  def test_load_latest_checkpoint(self):
    state_1 = {'w': jnp.ones(10)}
    state_2 = {'w': jnp.ones(10) * 2}
    checkpoint.save_checkpoint(self._root_dir, state_1, round_num=1, keep=2)
    checkpoint.save_checkpoint(self._root_dir, state_2, round_num=2, keep=2)


    latest_state, latest_round_num = checkpoint.load_latest_checkpoint(