# Path: vmoe/checkpoints/base.py
# CHECKPOINT_REGEX = re.compile(
#     r'^(.*?)(_[0-9]+)?(\.[a-zA-Z]+)?(-[0-9]+-of-[0-9]+)?$')
# def add_shard_suffix(filepath: str, shard: int, shard_count: int) -> str:
# def find_latest_complete_checkpoint_for_prefix(
#     prefix: str, suffixes: Optional[Sequence[str]] = None) -> Optional[str]:
# def iterate_complete_steps_for_prefix(
#     prefix: str,
#     suffixes: Optional[Sequence[str]] = None,
#     decreasing: bool = False) -> Iterator[int]:
#   def _parse_step_and_suffix_or_error(filepath):
# def remove_checkpoints(filepaths: Iterable[str],
#                        match_fn: Callable[[str], bool],
#                        *,
#                        thread_pool: Optional[ThreadPool] = None):
#   def remove(filepath):
# def remove_shard_suffix(filepath: str) -> str:
# def restore_checkpoint(filepath: str, tree: Optional[Any] = None) -> Any:
# def restore_multiple_checkpoints(
#     filepath_tree_map: Mapping[str, Any],
#     *,
#     thread_pool: Optional[ThreadPool] = None) -> Mapping[str, Any]:
# def save_checkpoint(filepath: str,
#                     tree: Any,
#                     *,
#                     overwrite: bool = True,
#                     makedirs: bool = True) -> str:
# def save_checkpoint_async(
#     filepath: str,
#     tree: Any,
#     *,
#     overwrite: bool = True,
#     makedirs: bool = True,
#     thread_pool: Optional[ThreadPool] = None) -> AsyncResult:
# def save_multiple_checkpoints_async(
#     filepath_tree_map: Mapping[str, Any],
#     *,
#     overwrite: bool = True,
#     makedirs: bool = True,
#     thread_pool: Optional[ThreadPool] = None) -> AsyncResult:


import os
import re
import jax
import jax.numpy as jnp
import numpy as np
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.checkpoints import base
            base.add_shard_suffix(prefix, shard, shard_count)), prefix)

  @parameterized.named_parameters(
      ('none', [], None),
      ('multiple', [4, 3], 'foo_4'),
  )
  @mock.patch.object(base, 'iterate_complete_steps_for_prefix')
  def test_find_latest_complete_checkpoint_for_prefix(
      self, mock_iterate_result, expected, mock_iterate):
    mock_iterate.return_value = mock_iterate_result
    self.assertEqual(
        base.find_latest_complete_checkpoint_for_prefix('foo'),
        expected)

  @parameterized.named_parameters(
      ('no_matches', [], [], None, False, []),
      # Return step 1 because 2 has tmp files.
      ('tmp_files_present',
       ['foo_1-0-of-1', 'foo_2'], ['.tmp.foo_2-0-of-1'], None, False, [1]),
      # Return step 1 because 2 is missing required suffix .bar.
      ('missing_suffix', ['foo_1.bar', 'foo_2'], [], ['.bar'], False, [1]),
      # Return step 1 and 2 in increasing/decreasing order.
      ('increasing', ['foo_1-0-of-1', 'foo_2-0-of-1'], [], [], False, [1, 2]),
      ('decreasing', ['foo_1-0-of-1', 'foo_2-0-of-1'], [], [], True, [2, 1]),
  )
  @mock.patch.object(base.gfile, 'glob')
  def test_iterate_complete_steps_for_prefix(
      self, files, tmp_files, suffixes, decreasing, expected, mock_glob):
    mock_glob.side_effect = [files, tmp_files]
    output = list(

# Path: vmoe/checkpoints/base.py
# CHECKPOINT_REGEX = re.compile(
#     r'^(.*?)(_[0-9]+)?(\.[a-zA-Z]+)?(-[0-9]+-of-[0-9]+)?$')
# def add_shard_suffix(filepath: str, shard: int, shard_count: int) -> str:
# def find_latest_complete_checkpoint_for_prefix(
#     prefix: str, suffixes: Optional[Sequence[str]] = None) -> Optional[str]:
# def iterate_complete_steps_for_prefix(
#     prefix: str,
#     suffixes: Optional[Sequence[str]] = None,
#     decreasing: bool = False) -> Iterator[int]:
#   def _parse_step_and_suffix_or_error(filepath):
# def remove_checkpoints(filepaths: Iterable[str],
#                        match_fn: Callable[[str], bool],
#                        *,
#                        thread_pool: Optional[ThreadPool] = None):
#   def remove(filepath):
# def remove_shard_suffix(filepath: str) -> str:
# def restore_checkpoint(filepath: str, tree: Optional[Any] = None) -> Any:
# def restore_multiple_checkpoints(
#     filepath_tree_map: Mapping[str, Any],
#     *,
#     thread_pool: Optional[ThreadPool] = None) -> Mapping[str, Any]:
# def save_checkpoint(filepath: str,
#                     tree: Any,
#                     *,
#                     overwrite: bool = True,
#                     makedirs: bool = True) -> str:
# def save_checkpoint_async(
#     filepath: str,
#     tree: Any,
#     *,
#     overwrite: bool = True,
#     makedirs: bool = True,
#     thread_pool: Optional[ThreadPool] = None) -> AsyncResult:
# def save_multiple_checkpoints_async(
#     filepath_tree_map: Mapping[str, Any],
#     *,
#     overwrite: bool = True,
#     makedirs: bool = True,
#     thread_pool: Optional[ThreadPool] = None) -> AsyncResult:


# Path: vmoe/checkpoints/base_test.py
import os
import re
import jax
import jax.numpy as jnp
import numpy as np
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.checkpoints import base

            base.add_shard_suffix(prefix, shard, shard_count)), prefix)

  @parameterized.named_parameters(
      ('none', [], None),
      ('multiple', [4, 3], 'foo_4'),
  )
  @mock.patch.object(base, 'iterate_complete_steps_for_prefix')
  def test_find_latest_complete_checkpoint_for_prefix(
      self, mock_iterate_result, expected, mock_iterate):
    mock_iterate.return_value = mock_iterate_result
    self.assertEqual(
        base.find_latest_complete_checkpoint_for_prefix('foo'),
        expected)

  @parameterized.named_parameters(
      ('no_matches', [], [], None, False, []),
      # Return step 1 because 2 has tmp files.
      ('tmp_files_present',
       ['foo_1-0-of-1', 'foo_2'], ['.tmp.foo_2-0-of-1'], None, False, [1]),
      # Return step 1 because 2 is missing required suffix .bar.
      ('missing_suffix', ['foo_1.bar', 'foo_2'], [], ['.bar'], False, [1]),
      # Return step 1 and 2 in increasing/decreasing order.
      ('increasing', ['foo_1-0-of-1', 'foo_2-0-of-1'], [], [], False, [1, 2]),
      ('decreasing', ['foo_1-0-of-1', 'foo_2-0-of-1'], [], [], True, [2, 1]),
  )
  @mock.patch.object(base.gfile, 'glob')
  def test_iterate_complete_steps_for_prefix(
      self, files, tmp_files, suffixes, decreasing, expected, mock_glob):
    mock_glob.side_effect = [files, tmp_files]
    output = list(

        base.iterate_complete_steps_for_prefix('foo', suffixes, decreasing))