diff --git a/src/clusterfuzz/_internal/bot/tasks/utasks/utask_utils.py b/src/clusterfuzz/_internal/base/task_utils.py similarity index 91% rename from src/clusterfuzz/_internal/bot/tasks/utasks/utask_utils.py rename to src/clusterfuzz/_internal/base/task_utils.py index f02414ec05..9f6f6f482f 100644 --- a/src/clusterfuzz/_internal/bot/tasks/utasks/utask_utils.py +++ b/src/clusterfuzz/_internal/base/task_utils.py @@ -1,4 +1,4 @@ -# Copyright 2023 Google LLC +# Copyright 2024 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """Module providing utilities for utask users. This module should not depend on -any other module in tasks to prevent circular imports.""" +any other module in tasks to prevent circular imports and issues with +appengine.""" from clusterfuzz._internal.system import environment diff --git a/src/clusterfuzz/_internal/base/tasks.py b/src/clusterfuzz/_internal/base/tasks.py index 6a230aeb2e..27453ffcf7 100644 --- a/src/clusterfuzz/_internal/base/tasks.py +++ b/src/clusterfuzz/_internal/base/tasks.py @@ -24,8 +24,8 @@ from clusterfuzz._internal.base import external_tasks from clusterfuzz._internal.base import persistent_cache +from clusterfuzz._internal.base import task_utils from clusterfuzz._internal.base import utils -from clusterfuzz._internal.bot.tasks.utasks import utask_utils from clusterfuzz._internal.config import local_config from clusterfuzz._internal.datastore import data_types from clusterfuzz._internal.datastore import ndb_utils @@ -271,7 +271,7 @@ def is_done_collecting_messages(): def get_postprocess_task(): """Gets a postprocess task if one exists.""" # This should only be run on non-preemptible bots. - if not utask_utils.is_remotely_executing_utasks(): + if not task_utils.is_remotely_executing_utasks(): return None pubsub_puller = PubSubPuller(POSTPROCESS_QUEUE) logs.log('Pulling from postprocess queue') @@ -449,7 +449,7 @@ def get_task_from_message(message) -> Optional[PubSubTask]: def get_utask_mains() -> List[PubSubTask]: """Returns a list of tasks for preprocessing many utasks on this bot and then running the uworker_mains in the same batch job.""" - if not utask_utils.is_remotely_executing_utasks(): + if not task_utils.is_remotely_executing_utasks(): return None pubsub_puller = PubSubPuller(UTASK_MAINS_QUEUE) messages = pubsub_puller.get_messages_time_limited(MAX_UTASKS, diff --git a/src/clusterfuzz/_internal/bot/tasks/setup.py b/src/clusterfuzz/_internal/bot/tasks/setup.py index 97fd5da297..7f2451cbe6 100644 --- a/src/clusterfuzz/_internal/bot/tasks/setup.py +++ b/src/clusterfuzz/_internal/bot/tasks/setup.py @@ -22,10 +22,10 @@ from clusterfuzz._internal.base import dates from clusterfuzz._internal.base import errors +from clusterfuzz._internal.base import task_utils from clusterfuzz._internal.base import tasks from clusterfuzz._internal.base import utils from clusterfuzz._internal.bot import testcase_manager -from clusterfuzz._internal.bot.tasks.utasks import utask_utils from clusterfuzz._internal.bot.tasks.utasks import uworker_io from clusterfuzz._internal.build_management import revisions from clusterfuzz._internal.datastore import data_handler @@ -189,7 +189,7 @@ def handle_setup_testcase_error(uworker_output: uworker_msg_pb2.Output): uworker_output.error_message) # Then reschedule the task. - command = utask_utils.get_command_from_module( + command = task_utils.get_command_from_module( uworker_output.uworker_input.module_name) testcase_fail_wait = environment.get_value('FAIL_WAIT') diff --git a/src/clusterfuzz/_internal/bot/tasks/task_types.py b/src/clusterfuzz/_internal/bot/tasks/task_types.py index 75c1f5889c..2d480cd857 100644 --- a/src/clusterfuzz/_internal/bot/tasks/task_types.py +++ b/src/clusterfuzz/_internal/bot/tasks/task_types.py @@ -14,8 +14,8 @@ """Types of tasks. This needs to be seperate from commands.py because base/tasks.py depends on this module and many things commands.py imports depend on base/tasks.py (i.e. avoiding circular imports).""" +from clusterfuzz._internal.base import task_utils from clusterfuzz._internal.bot.tasks import utasks -from clusterfuzz._internal.bot.tasks.utasks import utask_utils from clusterfuzz._internal.google_cloud_utils import batch from clusterfuzz._internal.metrics import logs from clusterfuzz._internal.system import environment @@ -101,7 +101,7 @@ class UTask(BaseUTask): @staticmethod def is_execution_remote(): - return utask_utils.is_remotely_executing_utasks() + return task_utils.is_remotely_executing_utasks() def execute(self, task_argument, job_type, uworker_env): """Executes a utask locally.""" diff --git a/src/clusterfuzz/_internal/google_cloud_utils/batch.py b/src/clusterfuzz/_internal/google_cloud_utils/batch.py index be12394029..c03154e806 100644 --- a/src/clusterfuzz/_internal/google_cloud_utils/batch.py +++ b/src/clusterfuzz/_internal/google_cloud_utils/batch.py @@ -20,8 +20,8 @@ from google.cloud import batch_v1 as batch from clusterfuzz._internal.base import retry +from clusterfuzz._internal.base import task_utils from clusterfuzz._internal.base import utils -from clusterfuzz._internal.bot.tasks.utasks import utask_utils from clusterfuzz._internal.config import local_config from clusterfuzz._internal.datastore import data_types from clusterfuzz._internal.metrics import logs @@ -91,7 +91,7 @@ def __init__(self, command, job_type, input_download_url): def create_uworker_main_batch_job(module, job_type, input_download_url): - command = utask_utils.get_command_from_module(module) + command = task_utils.get_command_from_module(module) batch_tasks = [BatchTask(command, job_type, input_download_url)] result = create_uworker_main_batch_jobs(batch_tasks) if result is None: diff --git a/src/clusterfuzz/_internal/tests/core/base/task_utils_test.py b/src/clusterfuzz/_internal/tests/core/base/task_utils_test.py new file mode 100644 index 0000000000..7a103e71d6 --- /dev/null +++ b/src/clusterfuzz/_internal/tests/core/base/task_utils_test.py @@ -0,0 +1,35 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for task_utils.""" +import unittest + +from clusterfuzz._internal.base import task_utils +from clusterfuzz._internal.bot.tasks import commands + + +class GetCommandFromModuleTest(unittest.TestCase): + """Tests for get_command_from_module.""" + + def test_get_command_from_module(self): + # pylint: disable=protected-access + """Tests that get_command_from_module returns the correct command.""" + for command, module in commands._COMMAND_MODULE_MAP.items(): + if command in {'postprocess', 'uworker_main'}: + continue + self.assertEqual(command, + task_utils.get_command_from_module(module.__name__)) + with self.assertRaises(ValueError): + task_utils.get_command_from_module('postprocess') + with self.assertRaises(ValueError): + task_utils.get_command_from_module('uworker_main') diff --git a/src/clusterfuzz/_internal/tests/core/bot/tasks/task_types_test.py b/src/clusterfuzz/_internal/tests/core/bot/tasks/task_types_test.py index 38aed706a4..0c065c26e9 100644 --- a/src/clusterfuzz/_internal/tests/core/bot/tasks/task_types_test.py +++ b/src/clusterfuzz/_internal/tests/core/bot/tasks/task_types_test.py @@ -28,7 +28,7 @@ def setUp(self): def test_is_remote(self): with mock.patch( - 'clusterfuzz._internal.bot.tasks.utasks.utask_utils.is_remotely_executing_utasks', + 'clusterfuzz._internal.base.task_utils.is_remotely_executing_utasks', return_value=True): self.assertTrue(task_types.UTaskCombined.is_execution_remote()) @@ -44,7 +44,7 @@ def test_mac(self): job_name = 'libfuzzer_mac_asan' with mock.patch( - 'clusterfuzz._internal.bot.tasks.utasks.utask_utils.is_remotely_executing_utasks', + 'clusterfuzz._internal.base.task_utils.is_remotely_executing_utasks', return_value=True): data_types.Job(name=job_name, platform='MAC').put() self.assertFalse(task_types.is_remote_utask('variant', job_name)) @@ -53,7 +53,7 @@ def test_linux(self): job_name = 'libfuzzer_linux_asan' with mock.patch( - 'clusterfuzz._internal.bot.tasks.utasks.utask_utils.is_remotely_executing_utasks', + 'clusterfuzz._internal.base.task_utils.is_remotely_executing_utasks', return_value=True): data_types.Job(name=job_name, platform='LINUX').put() self.assertTrue(task_types.is_remote_utask('variant', job_name)) @@ -62,7 +62,7 @@ def test_trusted(self): job_name = 'libfuzzer_linux_asan' with mock.patch( - 'clusterfuzz._internal.bot.tasks.utasks.utask_utils.is_remotely_executing_utasks', + 'clusterfuzz._internal.base.task_utils.is_remotely_executing_utasks', return_value=True): data_types.Job(name=job_name, platform='LINUX').put() self.assertFalse(task_types.is_remote_utask('impact', job_name))