# Path: vmoe/train/periodic_actions.py
# class SingleProcessPeriodicAction:
# class ReportProgress(periodic_actions.ReportProgress):
#   def __call__(self, *args, **kwargs) -> bool:
#   def __call__(self, step: int, t: Optional[float] = None, **kwargs) -> bool:
#   def _apply_extra(self, step: int, t: float, scalar_metrics: Dict[str, Any]):


import contextlib
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.train import periodic_actions
"""Tests for periodic_actions."""



class SingleProcessPeriodicActionTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ('_should_not_run', 0, False),
      ('_should_run', 1, True),
  )
  def test(self, process_index, expected_output):
    action1 = mock.MagicMock(periodic_actions.PeriodicAction)
    action1.return_value = True
    action2 = periodic_actions.SingleProcessPeriodicAction(
        periodic_action=action1, process_index=1)
    with mock.patch.object(periodic_actions.jax, 'process_index',
                           return_value=process_index):
      self.assertEqual(action2(step=3), expected_output)


class ReportProgressTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ('_without_memory_profile', False,
       [{'foo': 0.4, 'bar/foo': 0.04}, {'foo': 0.8, 'bar/foo': 0.08}]),
      ('_with_memory_profile', True,
       [{'foo': 0.4, 'bar/foo': 0.04, 'host_memory_mb': mock.ANY},
        {'foo': 0.8, 'bar/foo': 0.08, 'host_memory_mb': mock.ANY}]),
  )
  def test(self, available, expected):

# Path: vmoe/train/periodic_actions.py
# class SingleProcessPeriodicAction:
# class ReportProgress(periodic_actions.ReportProgress):
#   def __call__(self, *args, **kwargs) -> bool:
#   def __call__(self, step: int, t: Optional[float] = None, **kwargs) -> bool:
#   def _apply_extra(self, step: int, t: float, scalar_metrics: Dict[str, Any]):


# Path: vmoe/train/periodic_actions_test.py
import contextlib
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.train import periodic_actions

"""Tests for periodic_actions."""



class SingleProcessPeriodicActionTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ('_should_not_run', 0, False),
      ('_should_run', 1, True),
  )
  def test(self, process_index, expected_output):
    action1 = mock.MagicMock(periodic_actions.PeriodicAction)
    action1.return_value = True
    action2 = periodic_actions.SingleProcessPeriodicAction(
        periodic_action=action1, process_index=1)
    with mock.patch.object(periodic_actions.jax, 'process_index',
                           return_value=process_index):
      self.assertEqual(action2(step=3), expected_output)


class ReportProgressTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ('_without_memory_profile', False,
       [{'foo': 0.4, 'bar/foo': 0.04}, {'foo': 0.8, 'bar/foo': 0.08}]),
      ('_with_memory_profile', True,
       [{'foo': 0.4, 'bar/foo': 0.04, 'host_memory_mb': mock.ANY},
        {'foo': 0.8, 'bar/foo': 0.08, 'host_memory_mb': mock.ANY}]),
  )
  def test(self, available, expected):

    writer = mock.MagicMock(periodic_actions.periodic_actions.MetricWriter)