mmengine/tests/test_hook/test_sampler_seed_hook.py
Mashiro a7961407e4
[Refactor] Refactor the interfaces of Hook and its subclassed (#117)
* Fix hook

* Fix

* Fix docs

* FIx

* Fix

* Fix as comment

* update

* Fix hook

* Fix hook

* Fix hook

* Fix itertimerhook

* Fix iter_timer_hook

* Fix

* Fix

* fix logger hook

* Fix loggerhook

* update cur_dataloader

* Fix docstring

* Fix docstring

* Fix as commet

* Fix as commet

* Fix as comment

* rename is_last_epoch, enhance and add after_val before_val .etc

* fix typo in docstring

* remove resolved TODO

* refactor docstring
2022-03-13 16:48:09 +08:00

29 lines
1014 B
Python

# Copyright (c) OpenMMLab. All rights reserved.
from unittest.mock import Mock
from mmengine.hooks import DistSamplerSeedHook
class TestDistSamplerSeedHook:
def test_before_epoch(self):
hook = DistSamplerSeedHook()
# Test dataset sampler
runner = Mock()
runner.epoch = 1
runner.cur_dataloader = Mock()
runner.cur_dataloader.sampler = Mock()
runner.cur_dataloader.sampler.set_epoch = Mock()
hook.before_train_epoch(runner)
runner.cur_dataloader.sampler.set_epoch.assert_called()
# Test batch sampler
runner = Mock()
runner.cur_dataloader = Mock()
runner.cur_dataloader.sampler = Mock(spec_set=True)
runner.cur_dataloader.batch_sampler = Mock()
runner.cur_dataloader.batch_sampler.sampler = Mock()
runner.cur_dataloader.batch_sampler.sampler.set_epoch = Mock()
hook.before_train_epoch(runner)
runner.cur_dataloader.batch_sampler.sampler.set_epoch.assert_called()