mmengine/tests/test_hook/test_iter_timer_hook.py
RangiLyu 9a61b389e7
[Refactor] Add batch_idx to hook input. (#140)
* [Refactor] Add batch_idx to hook input.

* update
2022-03-29 11:40:38 +08:00

30 lines
809 B
Python

# Copyright (c) OpenMMLab. All rights reserved.
from unittest.mock import Mock
from mmengine.hooks import IterTimerHook
class TestIterTimerHook:
def test_before_epoch(self):
hook = IterTimerHook()
runner = Mock()
hook._before_epoch(runner)
assert isinstance(hook.t, float)
def test_before_iter(self):
hook = IterTimerHook()
runner = Mock()
runner.log_buffer = dict()
hook._before_epoch(runner)
hook._before_iter(runner, 0)
runner.message_hub.update_log.assert_called()
def test_after_iter(self):
hook = IterTimerHook()
runner = Mock()
runner.log_buffer = dict()
hook._before_epoch(runner)
hook._after_iter(runner, 0)
runner.message_hub.update_log.assert_called()