diff --git a/gokart/redis_lock.py b/gokart/redis_lock.py index 6201b006..b405e1ad 100644 --- a/gokart/redis_lock.py +++ b/gokart/redis_lock.py @@ -94,7 +94,7 @@ def wrapper(*args, **kwargs): return wrapper -def wrap_with_run_lock(func, redis_params: RedisParams): +def wrap_with_run_lock(func: Callable, redis_params: RedisParams): """Redis lock wrapper function for RunWithLock. When a fucntion is wrapped by RunWithLock, the wrapped function will be simply wrapped with redis lock. https://github.com/m3dev/gokart/issues/265 diff --git a/gokart/run_with_lock.py b/gokart/run_with_lock.py index e3c6bdcd..6596423f 100644 --- a/gokart/run_with_lock.py +++ b/gokart/run_with_lock.py @@ -22,5 +22,5 @@ def _run_with_lock(cls, func, output_list: list): return func() output = output_list.pop() - wrapped_func = output.wrap_with_lock(func) + wrapped_func = output.wrap_with_run_lock(func) return cls._run_with_lock(func=wrapped_func, output_list=output_list) diff --git a/gokart/target.py b/gokart/target.py index b920af87..7dc900c8 100644 --- a/gokart/target.py +++ b/gokart/target.py @@ -44,7 +44,7 @@ def last_modification_time(self) -> datetime: def path(self) -> str: return self._path() - def wrap_with_lock(self, func): + def wrap_with_run_lock(self, func): return wrap_with_run_lock(func=func, redis_params=self._get_redis_params()) @abstractmethod diff --git a/test/test_redis_lock.py b/test/test_redis_lock.py index 081ea80e..6b4a3afb 100644 --- a/test/test_redis_lock.py +++ b/test/test_redis_lock.py @@ -5,7 +5,8 @@ import fakeredis -from gokart.redis_lock import RedisClient, RedisParams, make_redis_key, make_redis_params, wrap_with_dump_lock, wrap_with_remove_lock, wrap_with_run_lock +from gokart.redis_lock import (RedisClient, RedisParams, make_redis_key, make_redis_params, wrap_with_dump_lock, wrap_with_load_lock, wrap_with_remove_lock, + wrap_with_run_lock) class TestRedisClient(unittest.TestCase): @@ -253,7 +254,7 @@ def test_no_redis(self): redis_port=None, ) mock_func = MagicMock() - resulted = wrap_with_run_lock(func=mock_func, redis_params=redis_params)(123, b='abc') + resulted = wrap_with_load_lock(func=mock_func, redis_params=redis_params)(123, b='abc') mock_func.assert_called_once() called_args, called_kwargs = mock_func.call_args @@ -273,7 +274,7 @@ def test_use_redis(self): with patch('gokart.redis_lock.redis.Redis') as redis_mock: redis_mock.side_effect = fakeredis.FakeRedis mock_func = MagicMock() - resulted = wrap_with_run_lock(func=mock_func, redis_params=redis_params)(123, b='abc') + resulted = wrap_with_load_lock(func=mock_func, redis_params=redis_params)(123, b='abc') mock_func.assert_called_once() called_args, called_kwargs = mock_func.call_args @@ -294,7 +295,7 @@ def test_check_lock_extended(self): with patch('gokart.redis_lock.redis.Redis') as redis_mock: redis_mock.side_effect = fakeredis.FakeRedis - resulted = wrap_with_run_lock(func=_sample_long_func, redis_params=redis_params)(123, b='abc') + resulted = wrap_with_load_lock(func=_sample_long_func, redis_params=redis_params)(123, b='abc') expected = dict(a=123, b='abc') self.assertEqual(resulted, expected) @@ -311,7 +312,7 @@ def test_lock_is_removed_after_func_is_finished(self): with patch('gokart.redis_lock.redis.Redis') as redis_mock: redis_mock.return_value = fakeredis.FakeRedis(server=server, host=redis_params.redis_host, port=redis_params.redis_port) mock_func = MagicMock() - resulted = wrap_with_run_lock(func=mock_func, redis_params=redis_params)(123, b='abc') + resulted = wrap_with_load_lock(func=mock_func, redis_params=redis_params)(123, b='abc') mock_func.assert_called_once() called_args, called_kwargs = mock_func.call_args @@ -336,7 +337,7 @@ def test_lock_is_removed_after_func_is_finished_with_error(self): with patch('gokart.redis_lock.redis.Redis') as redis_mock: redis_mock.return_value = fakeredis.FakeRedis(server=server, host=redis_params.redis_host, port=redis_params.redis_port) try: - wrap_with_run_lock(func=_sample_func_with_error, redis_params=redis_params)(123, b='abc') + wrap_with_load_lock(func=_sample_func_with_error, redis_params=redis_params)(123, b='abc') except Exception: fake_redis = fakeredis.FakeStrictRedis(server=server) with self.assertRaises(KeyError): diff --git a/test/test_task_on_kart.py b/test/test_task_on_kart.py index e31c5d73..75b462ba 100644 --- a/test/test_task_on_kart.py +++ b/test/test_task_on_kart.py @@ -522,7 +522,7 @@ def test_run_with_lock_decorator(self): def _wrap(func): return func - with patch('gokart.target.TargetOnKart.wrap_with_lock') as mock_obj: + with patch('gokart.target.TargetOnKart.wrap_with_run_lock') as mock_obj: mock_obj.side_effect = _wrap task.run() mock_obj.assert_called_once() @@ -533,7 +533,7 @@ def test_run_with_lock_decorator_multiple_output(self): def _wrap(func): return func - with patch('gokart.target.TargetOnKart.wrap_with_lock') as mock_obj: + with patch('gokart.target.TargetOnKart.wrap_with_run_lock') as mock_obj: mock_obj.side_effect = _wrap task.run() self.assertEqual(mock_obj.call_count, 2) @@ -544,7 +544,7 @@ def test_run_without_lock_decorator(self): def _wrap(func): return func - with patch('gokart.target.TargetOnKart.wrap_with_lock') as mock_obj: + with patch('gokart.target.TargetOnKart.wrap_with_run_lock') as mock_obj: mock_obj.side_effect = _wrap task.run() mock_obj.assert_not_called()