Separate creation and setup requests in OT creation process

framework/origin_trials_client.py

@@ -29,6 +29,7 @@
 
 
 class RequestTrial(TypedDict):
+  id: NotRequired[int]
   display_name: str
   start_milestone: str
   end_milestone: str
@@ -58,6 +59,12 @@ class CreateOriginTrialRequest(TypedDict):
   trial_contacts: list[str]
 
 
+class SetUpTrialRequest(TypedDict):
+  trial_id: int
+  announcement_groups_owners: list[str]
+  trial_contacts: list[str]
+
+
 def get_trials_list() -> list[dict[str, Any]]:
   """Get a list of all origin trials.
 
@@ -132,21 +139,14 @@ def _get_ot_access_token() -> str:
   return credentials.token
 
 
-def create_origin_trial(ot_stage: Stage) -> str | None:
-  """Create an origin trial.
-
-  Raises:
-    requests.exceptions.RequestException: If the request fails to connect or
-      the HTTP status code is not successful.
+def _send_create_trial_request(
+    ot_stage: Stage, api_key: str, access_token: str
+  ) -> tuple[int|None, str|None]:
+  """Send an origin trial creation request to the origin trials API.
+  Returns:
+    Newly created origin trial ID if trial was created, any error text if there
+    was an issue during the creation process.
   """
-  if settings.DEV_MODE:
-    logging.info('Creation request will not be sent to origin trials API in '
-                 'local environment.')
-    return None
-  key = secrets.get_ot_api_key()
-  if key is None:
-    return None
-
   json: CreateOriginTrialRequest = {
     'trial': {
       'display_name': ot_stage.ot_display_name,
@@ -167,16 +167,6 @@ def create_origin_trial(ot_stage: Stage) -> str | None:
     'registration_config': {'approval_type': 'NONE'},
     'trial_contacts': []
   }
-
-  # Get a list of all OT @google.com contacts (ot_owner_email must be a google
-  # contact).
-  unique_contacts = [ot_stage.ot_owner_email]
-  unique_contacts.extend(ot_stage.ot_emails)
-  unique_contacts = [email for email in set(unique_contacts)
-                     if email.endswith('@google.com')]
-  json['trial_contacts'] = unique_contacts
-  if ot_stage.ot_chromium_trial_name:
-    json['trial']['origin_trial_feature_name'] = ot_stage.ot_chromium_trial_name
   if ot_stage.ot_require_approvals:
     json['registration_config'] = {
       'approval_type': 'CUSTOM',
@@ -185,23 +175,99 @@ def create_origin_trial(ot_stage: Stage) -> str | None:
       'approval_group_email': ot_stage.ot_approval_group_email,
       'approval_buganizer_custom_field_id': ot_stage.ot_approval_buganizer_custom_field_id,
     }
+  if ot_stage.ot_chromium_trial_name:
+    json['trial']['origin_trial_feature_name'] = ot_stage.ot_chromium_trial_name
   if ot_stage.ot_is_deprecation_trial:
     json['registration_config']['allow_public_suffix_subdomains'] = True
-  access_token = _get_ot_access_token()
+
   headers = {'Authorization': f'Bearer {access_token}'}
-  url = f'{settings.OT_API_URL}/v1/trials:setup'
+  url = f'{settings.OT_API_URL}/v1/trials:initialize'
 
+  # Retry the request a number of times if any issues arise.
   try:
     response = requests.post(
-        url, headers=headers, params={'key': key}, json=json)
-    logging.info(response.text)
+        url, headers=headers, params={'key': api_key}, json=json)
+    logging.info(f'CreateTrial response text: {response.text}')
     response.raise_for_status()
-  except requests.exceptions.RequestException as e:
+  except requests.exceptions.RequestException:
     logging.exception(
         f'Failed to get response from origin trials API. {response.text}')
-    raise e
+    return None, response.text
+  response_json = response.json()
+  return response_json['trial']['id'], None
+
+
+def _send_set_up_trial_request(
+    trial_id: int,
+    owners: list[str],
+    contacts: list[str],
+    api_key: str,
+    access_token: str
+  ) -> str|None:
+  """Send an origin trial setup request to the origin trials API.
+  Returns:
+    Any error text if there was an issue during the setup process.
+  """
+  json: SetUpTrialRequest = {
+    'trial_id': trial_id,
+    'announcement_groups_owners': owners,
+    'trial_contacts': contacts,
+  }
+  headers = {'Authorization': f'Bearer {access_token}'}
+  url = f'{settings.OT_API_URL}/v1/trials/{trial_id}:setup'
+  try:
+    response = requests.post(
+        url, headers=headers, params={'key': api_key}, json=json)
+    logging.info(f'SetUpTrial response text: {response.text}')
+    response.raise_for_status()
+  except requests.exceptions.RequestException:
+    logging.exception(
+        f'Failed to get response from origin trials API. {response.text}')
+    return response.text
+  return None
+
+
+def create_origin_trial(ot_stage: Stage) -> tuple[str|None, str|None]:
+  """Send an origin trial creation request to the origin trials API.
+
+  Raises:
+    requests.exceptions.RequestException: If the request fails to connect or
+      the HTTP status code is not successful.
+  Returns:
+    Newly created origin trial ID if trial was created, any error text if there
+    was an issue during the creation process.
+  """
+  if settings.DEV_MODE:
+    logging.info('Creation request will not be sent to origin trials API in '
+                 'local environment.')
+    return None, None
+  key = secrets.get_ot_api_key()
+  if key is None:
+    return None, 'No API key found for origin trials API'
+
+  # Get a list of all OT @google.com contacts (ot_owner_email must be a google
+  # contact).
+  unique_contacts = [ot_stage.ot_owner_email]
+  unique_contacts.extend(ot_stage.ot_emails)
+  unique_contacts = [email for email in set(unique_contacts)
+                     if email.endswith('@google.com')]
+  # A trial must have a google.com domain contact.
+  if len(unique_contacts) == 0:
+    return None, 'No trial contacts found in google.com domain'
+
+  access_token = _get_ot_access_token()
+  origin_trial_id, error_text = _send_create_trial_request(
+      ot_stage, key, access_token)
+
+  if origin_trial_id is None:
+    return None, error_text
+
+  error_text = _send_set_up_trial_request(
+      # TODO(DanielRyanSmith): Add owners contacts in subsequent PR.
+      origin_trial_id, [], unique_contacts, key, access_token)
+
+  return str(origin_trial_id), error_text
 
-  return str(response.json()['id'])
 
 def activate_origin_trial(origin_trial_id: str) -> None:
   """Activate an existing origin trial.

framework/origin_trials_client_test.py

@@ -189,8 +189,10 @@ def test_create_origin_trial__no_api_key(
       self, mock_requests_post, mock_api_key_get):
     """If no API key is available, do not send creation request."""
     mock_api_key_get.return_value = None
-    origin_trials_client.create_origin_trial(self.ot_stage)
+    ot_id, error_text = origin_trials_client.create_origin_trial(self.ot_stage)
 
+    self.assertIsNone(ot_id)
+    self.assertEqual('No API key found for origin trials API', error_text)
     mock_api_key_get.assert_called_once()
     # POST request should not be executed with no API key.
     mock_requests_post.assert_not_called()
@@ -204,22 +206,21 @@ def test_create_origin_trial__with_api_key(
       mock_get_ot_access_token, mock_api_key_get):
     """If an API key is available, POST should create trial and return true."""
     mock_requests_post.return_value = mock.MagicMock(
-        status_code=200, json=lambda : {'id': -1234567890})
+        status_code=200, json=lambda : (
+            {'trial': {'id': -1234567890}, 'should_retry': False}))
     mock_get_trial_end_time.return_value = 111222333
     mock_get_ot_access_token.return_value = 'access_token'
     mock_api_key_get.return_value = 'api_key_value'
 
-    id = origin_trials_client.create_origin_trial(self.ot_stage)
-    self.assertEqual(id, '-1234567890')
+    ot_id, error_text = origin_trials_client.create_origin_trial(self.ot_stage)
+    self.assertEqual(ot_id, '-1234567890')
+    self.assertIsNone(error_text)
 
     mock_api_key_get.assert_called_once()
     mock_get_ot_access_token.assert_called_once()
-    mock_requests_post.assert_called_once()
-    # unordered list is in the request, so compare args individually.
-    json_args = mock_requests_post.call_args_list[0][1]['json']
-    # Only unique @google.com emails should be sent as contacts.
-    self.assertCountEqual(['[email protected]', '[email protected]'],
-                          json_args['trial_contacts'])
+    # Two separate POST requests made.
+    self.assertEqual(2, mock_requests_post.call_count)
+    create_trial_json = mock_requests_post.call_args_list[0][1]['json']
     self.assertEqual({
           'display_name': 'Example Trial',
           'start_milestone': '100',
@@ -234,16 +235,21 @@ def test_create_origin_trial__with_api_key(
           'chromestatus_url': f'{settings.SITE_URL}feature/1',
           'allow_third_party_origins': True,
           'type': 'DEPRECATION',
-        }, json_args['trial'])
+        }, create_trial_json['trial'])
     self.assertEqual({
           'allow_public_suffix_subdomains': True,
           'approval_type': 'CUSTOM',
           'approval_buganizer_component_id': 123456,
           'approval_criteria_url': 'https://example.com/criteria',
           'approval_group_email': '[email protected]',
           'approval_buganizer_custom_field_id': 111111,
-        }, json_args['registration_config'])
+        }, create_trial_json['registration_config'])
 
+    set_up_trial_json = mock_requests_post.call_args_list[1][1]['json']
+    # Only unique @google.com emails should be sent as contacts.
+    self.assertCountEqual(['[email protected]', '[email protected]'],
+                          set_up_trial_json['trial_contacts'])
+    self.assertEqual(-1234567890, set_up_trial_json['trial_id'])
 
   @mock.patch('framework.secrets.get_ot_api_key')
   @mock.patch('requests.post')

internals/maintenance_scripts.py

@@ -26,7 +26,6 @@
 from framework import utils
 from internals import approval_defs
 from internals.core_models import FeatureEntry, MilestoneSet, Stage
-from internals.data_types import StageDict
 from internals.review_models import Gate, Vote, Activity
 from internals.core_enums import *
 from internals.feature_links import batch_index_feature_entries
@@ -523,57 +522,70 @@ def get_template_data(self, **kwargs) -> str:
 
 class CreateOriginTrials(FlaskHandler):
 
-  def handle_creation(self, stage: Stage, stage_dict: StageDict) -> str | None:
+  def _send_creation_result_notification(
+      self, task_path: str, stage: Stage, params: dict|None = None) -> None:
+    if not params:
+      params = {}
+    print('sending email task to', task_path, 'with params', params)
+    stage_dict = converters.stage_to_json_dict(stage)
+    params['stage'] = stage_dict
+    cloud_tasks_helpers.enqueue_task(task_path, params)
+
+  def handle_creation(self, stage: Stage) -> bool:
     """Send a flagged creation request for processing to the Origin Trials
     API.
     """
-    new_id = None
-    # TODO(DanielRyanSmith): This request should have a retry process based on
-    # the error code returned from the OT server.
-    try:
-      new_id = origin_trials_client.create_origin_trial(stage)
-      stage.ot_setup_status = OT_CREATED
-    except requests.RequestException:
+    origin_trial_id, error_text = origin_trials_client.create_origin_trial(
+        stage)
+    if origin_trial_id:
+      stage.origin_trial_id = origin_trial_id
+    if error_text:
       logging.warning('Origin trial could not be created for stage '
-                      f'{stage.key.integer_id()}')
-      cloud_tasks_helpers.enqueue_task(
-          '/tasks/email-ot-creation-request-failed', {'stage': stage_dict})
+                     f'{stage.key.integer_id()}')
       stage.ot_setup_status = OT_CREATION_FAILED
-    stage.put()
-    return new_id
+      self._send_creation_result_notification(
+          '/tasks/email-ot-creation-request-failed',
+          stage,
+          {'error_text': error_text})
+      return False
+    else:
+      stage.ot_setup_status = OT_CREATED
+      logging.info(f'Origin trial created for stage {stage.key.integer_id()}')
+    return True
 
-  def handle_activation(self, stage: Stage, stage_dict: StageDict) -> None:
+  def handle_activation(self, stage: Stage) -> None:
     """Send trial activation request."""
     try:
       origin_trials_client.activate_origin_trial(stage.origin_trial_id)
-      cloud_tasks_helpers.enqueue_task(
-          '/tasks/email-ot-activated', {'stage': stage_dict})
       stage.ot_setup_status = OT_ACTIVATED
+      self._send_creation_result_notification(
+          '/tasks/email-ot-activated', stage)
     except requests.RequestException:
-      cloud_tasks_helpers.enqueue_task(
-          '/tasks/email-ot-activation-failed', {'stage': stage_dict})
-      stage.ot_setup_status = OT_ACTIVATION_FAILED
       # The activation still needs to occur,
       # so the activation date is set for current date.
       stage.ot_activation_date = date.today()
+      stage.ot_setup_status = OT_ACTIVATION_FAILED
+      self._send_creation_result_notification(
+          '/tasks/email-ot-activation-failed', stage)
 
   def _get_today(self) -> date:
     return date.today()
 
-  def prepare_for_activation(self, stage: Stage, stage_dict: StageDict) -> None:
+  def prepare_for_activation(self, stage: Stage) -> None:
     """Set up activation date or activate trial now."""
     mstone_info = utils.get_chromium_milestone_info(
         stage.milestones.desktop_first)
     date = datetime.strptime(
         mstone_info['mstones'][0]['branch_point'],
         utils.CHROMIUM_SCHEDULE_DATE_FORMAT).date()
     if date <= self._get_today():
-      self.handle_activation(stage, stage_dict)
+      print('sending for activation. Today:', self._get_today(), 'branch_date: ', date)
+      self.handle_activation(stage)
     else:
       stage.ot_activation_date = date
       stage.ot_setup_status = OT_CREATED
-      cloud_tasks_helpers.enqueue_task(
-          '/tasks/email-ot-creation-processed', {'stage': stage_dict})
+      self._send_creation_result_notification(
+          '/tasks/email-ot-creation-processed', stage)
 
   def get_template_data(self, **kwargs) -> str:
     """Create any origin trials that are flagged for creation."""
@@ -586,11 +598,9 @@ def get_template_data(self, **kwargs) -> str:
         Stage.ot_setup_status == OT_READY_FOR_CREATION).fetch()
     for stage in ot_stages:
       stage.ot_action_requested = False
-      stage_dict = converters.stage_to_json_dict(stage)
-      origin_trial_id = self.handle_creation(stage, stage_dict)
-      if origin_trial_id:
-        stage.origin_trial_id = origin_trial_id
-        self.prepare_for_activation(stage, stage_dict)
+      creation_success = self.handle_creation(stage)
+      if creation_success:
+        self.prepare_for_activation(stage)
       stage.put()
 
     return f'{len(ot_stages)} trial creation request(s) processed.'