prepare unscheduled_tasks

Signed-off-by: Javan Lacerda <javanlacerda@google.com>

Author: javanlacerda

src/clusterfuzz/_internal/batch/service.py

@@ -375,8 +375,10 @@ def create_utask_main_job(self, module: str, job_type: str,
     batch_tasks = [
         remote_task_types.RemoteTask(command, job_type, input_download_url)
     ]
-    uncreated_tasks = self.create_utask_main_jobs(batch_tasks)
-    return uncreated_tasks
+    result = self.create_utask_main_jobs(batch_tasks)
+    if not result:
+      return None
+    return result[0]
 
   def create_utask_main_jobs(self,
                              remote_tasks: List[remote_task_types.RemoteTask]):

src/clusterfuzz/_internal/remote_task/remote_task_gate.py

@@ -21,8 +21,10 @@
 
 import collections
 import random
-from typing import List
+import typing
 
+from clusterfuzz._internal.datastore import data_types
+from clusterfuzz._internal.datastore import ndb_utils
 from clusterfuzz._internal.metrics import logs
 from clusterfuzz._internal.remote_task import remote_task_adapters
 from clusterfuzz._internal.remote_task import remote_task_types
@@ -105,13 +107,39 @@ def get_job_frequency(self):
     logs.info('Job frequencies', frequencies=frequencies)
     return frequencies
 
+  def prepare_unscheduled_tasks(
+      self, unscheduled_remote_tasks: typing.List[remote_task_types.RemoteTask]
+  ) -> typing.List[remote_task_types.RemoteTask]:
+    """Prepares the unscheduled remote tasks to be sent back to the 
+      preprocess queue.
+
+      The messages in the preprocess queue expects tasks with 
+      the properties followng task_name, fuzzer, job, eta. 
+      The utasks has the signed url as argument instead of the fuzzer,
+      then for sending it back to the preprocess, it recovers its fuzzer
+      and override the argument.
+    """
+
+    job_names = {task.job_type for task in unscheduled_remote_tasks}
+    query = data_types.FuzzerJob.query(
+        data_types.FuzzerJob.job.IN(list(job_names)))
+    fuzzer_jobs = ndb_utils.get_all_from_query(query)
+    fuzzer_jobs_name_mapped = {job.name: job for job in fuzzer_jobs}
+    for task in unscheduled_remote_tasks:
+      if task.job_type not in fuzzer_jobs_name_mapped:
+        logs.error(f'{task.job_type} not found.')
+        continue
+      task.argument = fuzzer_jobs_name_mapped[task.job_type].argument
+
+    return unscheduled_remote_tasks
+
   def create_utask_main_job(self, module, job_type, input_download_url):
     adapter_id = self._get_adapter()
     service = self._service_map[adapter_id]
     return service.create_utask_main_job(module, job_type, input_download_url)
 
-  def create_utask_main_jobs(self,
-                             remote_tasks: List[remote_task_types.RemoteTask]):
+  def create_utask_main_jobs(
+      self, remote_tasks: typing.List[remote_task_types.RemoteTask]):
     """Creates a batch of remote tasks, distributing them across backends.
 
     This method handles two cases:
@@ -146,13 +174,16 @@ def create_utask_main_jobs(self,
         adapter_id = list(frequencies.keys())[i % len(frequencies)]
         tasks_by_adapter[adapter_id].append(task)
 
-    results = []
+    unscheduled_tasks = []
     for adapter_id, tasks in tasks_by_adapter.items():
       if tasks:
         try:
           logs.info(f'Sending {len(tasks)} tasks to {adapter_id}.')
           service = self._service_map[adapter_id]
-          results.extend(service.create_utask_main_jobs(tasks))
+          unscheduled_tasks.extend(service.create_utask_main_jobs(tasks))
         except Exception:  # pylint: disable=broad-except
           logs.error(f'Failed to send {len(tasks)} tasks to {adapter_id}.')
-    return results
+
+    prepared_unscheduled_tasks = self.prepare_unscheduled_tasks(
+        unscheduled_tasks)
+    return prepared_unscheduled_tasks

src/clusterfuzz/_internal/tests/core/batch/batch_service_test.py

@@ -254,7 +254,7 @@ def test_create_uworker_main_batch_job(self):
           UUIDS[0], spec1, ['url1'])
       self.mock_batch_client_instance.create_job.assert_called_with(
           expected_create_request)
-      self.assertEqual(result, [])
+      self.assertIsNone(result)
 
 
 @test_utils.with_cloud_emulators('datastore')

src/clusterfuzz/_internal/tests/core/remote_task/remote_task_gate_test.py

@@ -55,6 +55,12 @@ def setUp(self):
     self.patcher.start()
     self.addCleanup(self.patcher.stop)
 
+    self.mock_prepare_unscheduled_tasks = mock.patch.object(
+        remote_task_gate.RemoteTaskGate,
+        'prepare_unscheduled_tasks',
+        side_effect=lambda x: x).start()
+    self.addCleanup(mock.patch.stopall)
+
   def test_init(self):
     """Tests that the RemoteTaskGate initializes correctly and creates
     service map."""
@@ -197,3 +203,33 @@ def test_create_utask_main_jobs_full_gcp_batch(self, mock_get_job_frequency):
     self.mock_gcp_batch_service.create_utask_main_jobs.assert_called_once_with(
         tasks)
     self.mock_k8s_service.create_utask_main_jobs.assert_not_called()
+
+  @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency')
+  def test_create_utask_main_jobs_calls_prepare_unscheduled_tasks(
+      self, mock_get_job_frequency):
+    """Tests that create_utask_main_jobs calls prepare_unscheduled_tasks with
+    unscheduled tasks."""
+    tasks = [
+        remote_task_types.RemoteTask('c', 'j', 'u1'),
+    ]
+    unscheduled_tasks = [
+        remote_task_types.RemoteTask('c', 'j', 'u1'),
+    ]
+    prepared_tasks = [
+        remote_task_types.RemoteTask('c', 'j', 'u1'),
+    ]
+    prepared_tasks[0].argument = 'arg'
+
+    mock_get_job_frequency.return_value = {'kubernetes': 1.0, 'gcp_batch': 0.0}
+    self.mock_k8s_service.create_utask_main_jobs.return_value = unscheduled_tasks
+
+    self.mock_prepare_unscheduled_tasks.side_effect = None
+    self.mock_prepare_unscheduled_tasks.return_value = prepared_tasks
+
+    gate = remote_task_gate.RemoteTaskGate()
+    result = gate.create_utask_main_jobs(tasks)
+
+    self.mock_k8s_service.create_utask_main_jobs.assert_called_once_with(tasks)
+    self.mock_prepare_unscheduled_tasks.assert_called_once_with(
+        unscheduled_tasks)
+    self.assertEqual(result, prepared_tasks)

src/clusterfuzz/_internal/tests/core/remote_task/remote_task_test.py

@@ -41,23 +41,19 @@ def setUp(self):
     mock_creds.token = 'fake-token'
     self.mock_auth_default.return_value = (mock_creds, 'test-project')
 
-    # Mock discovery.build to avoid network calls during KubernetesService init
-    patcher = mock.patch('googleapiclient.discovery.build')
+    # Mock _load_gke_credentials to avoid BOT_DIR dependency
+    patcher = mock.patch.object(k8s_service.KubernetesService,
+                                '_load_gke_credentials')
     self.addCleanup(patcher.stop)
-    self.mock_discovery_build = patcher.start()
-    mock_service = mock.Mock()
-    self.mock_discovery_build.return_value = mock_service
-    mock_service.projects().locations().clusters().list(
-    ).execute.return_value = {
-        'clusters': [{
-            'name': 'clusterfuzz-cronjobs-gke',
-            'endpoint': '1.2.3.4',
-            'masterAuth': {
-                'clusterCaCertificate':
-                    'ZmFrZS1jZXJ0'  # base64 encoded 'fake-cert'
-            }
-        }]
-    }
+    patcher.start()
+
+    # Mock prepare_unscheduled_tasks to avoid NDB context issues
+    patcher = mock.patch.object(
+        remote_task_gate.RemoteTaskGate,
+        'prepare_unscheduled_tasks',
+        side_effect=lambda x: x)
+    self.addCleanup(patcher.stop)
+    patcher.start()
 
     self.gate = remote_task_gate.RemoteTaskGate()