the-algorithm-ml/ml_logging/test_torch_logging.py

import unittest

from tml.ml_logging.torch_logging import logging


class Testlogging(unittest.TestCase):
  def test_warn_once(self):
    with self.assertLogs(level="INFO") as captured_logs:
      logging.info("first info")
      logging.warning("first warning")
      logging.warning("first warning")
      logging.info("second info")

    self.assertEqual(
      captured_logs.output,
      [
        "INFO:absl:first info",
        "WARNING:absl:first warning",
        "INFO:absl:second info",
      ],
    )