Origin trial creation cron job

cron.yaml

@@ -38,3 +38,8 @@ cron:
 - description: Send origin trial process reminder emails.
   url: /cron/send-ot-process-reminders
   schedule: every monday 4:00
+
+# TODO(DanielRyanSmith): Add this job when OT creation is fully implemented.
+# - description: Check if any origin trials require creation
+#   url: /cron/create_origin_trials
+#   schedule: every 5 minutes

framework/origin_trials_client.py

@@ -22,13 +22,12 @@
 import requests
 
 from framework import secrets
+from framework import utils
+from internals.core_models import Stage
 from internals.data_types import OriginTrialInfo
 import settings
 
 
-CHROMIUM_SCHEDULE_DATE_FORMAT = '%Y-%m-%dT%H:%M:%S'
-
-
 def get_trials_list() -> list[dict[str, Any]]:
   """Get a list of all origin trials.
 
@@ -75,24 +74,16 @@ def _get_trial_end_time(end_milestone: int) -> int:
       format.
   """
   milestone_plus_two = int(end_milestone) + 2
-  try:
-    response = requests.get(
-      'https://chromiumdash.appspot.com/fetch_milestone_schedule'
-      f'?mstone={milestone_plus_two}')
-    response.raise_for_status()
-  except requests.exceptions.RequestException as e:
-    logging.exception('Failed to get response from Chromium schedule API.')
-    raise e
-  response_json = response.json()
+  mstone_info = utils.get_chromium_milestone_info(milestone_plus_two)
 
   # Raise error if the response is not in the expected format.
-  if ('mstones' not in response_json
-      or len(response_json['mstones']) == 0
-      or 'late_stable_date' not in response_json['mstones'][0]):
+  if ('mstones' not in mstone_info
+      or len(mstone_info['mstones']) == 0
+      or 'late_stable_date' not in mstone_info['mstones'][0]):
     raise KeyError('Chromium schedule response not in expected format.')
   date = datetime.strptime(
-      response_json['mstones'][0]['late_stable_date'],
-      CHROMIUM_SCHEDULE_DATE_FORMAT)
+      mstone_info['mstones'][0]['late_stable_date'],
+      utils.CHROMIUM_SCHEDULE_DATE_FORMAT)
   return int(date.replace(tzinfo=timezone.utc).timestamp())
 
 
@@ -111,6 +102,88 @@ 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.
+  """
+  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 = {
+    'trial': {
+      'display_name': ot_stage.ot_display_name,
+      'start_milestone': str(ot_stage.milestones.desktop_first),
+      'end_milestone': str(ot_stage.milestones.desktop_last),
+      'end_time': {
+        'seconds': _get_trial_end_time(ot_stage.milestones.desktop_last)
+      },
+      'description': ot_stage.ot_description,
+      'documentation_url': ot_stage.ot_documentation_url,
+      'feedback_url': ot_stage.ot_feedback_submission_url,
+      'intent_to_experiment_url': ot_stage.intent_thread_url,
+      'chromestatus_url': f'{settings.SITE_URL}feature/{ot_stage.feature_id}',
+      'allow_third_party_origins': ot_stage.ot_has_third_party_support,
+      'type': ('DEPRECATION'
+                if ot_stage.ot_is_deprecation_trial else 'ORIGIN_TRIAL'),
+    }
+  }
+  if ot_stage.ot_chromium_trial_name:
+    json['origin_trial_feature_name'] = ot_stage.ot_chromium_trial_name
+  access_token = _get_ot_access_token()
+  headers = {'Authorization': f'Bearer {access_token}'}
+  url = f'{settings.OT_API_URL}/v1/trials-integration'
+
+  try:
+    response = requests.post(
+        url, headers=headers, params={'key': key}, json=json)
+    logging.info(response.text)
+    response.raise_for_status()
+  except requests.exceptions.RequestException as e:
+    logging.exception(
+        f'Failed to get response from origin trials API. {response.text}')
+    raise e
+
+  return response.json()['id']
+
+def activate_origin_trial(origin_trial_id: str) -> None:
+  """Activate an existing origin trial.
+
+  Raises:
+    requests.exceptions.RequestException: If the request fails to connect or
+      the HTTP status code is not successful.
+  """
+  if settings.DEV_MODE:
+    logging.info('Activation 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 = {'id': origin_trial_id}
+  access_token = _get_ot_access_token()
+  headers = {'Authorization': f'Bearer {access_token}'}
+  url = (f'{settings.OT_API_URL}/v1/trials-integration/'
+         f'{origin_trial_id}:activate')
+  try:
+    response = requests.post(
+        url, headers=headers, params={'key': key}, json=json)
+    logging.info(response.text)
+    response.raise_for_status()
+  except requests.exceptions.RequestException as e:
+    logging.exception(
+        f'Failed to get response from origin trials API. {response.text}')
+    raise e
+
+
 def extend_origin_trial(trial_id: str, end_milestone: int, intent_url: str):
   """Extend an existing origin trial.
 

framework/origin_trials_client_test.py

@@ -18,49 +18,65 @@
 from unittest import mock
 
 from framework import origin_trials_client
+from internals.core_models import MilestoneSet, Stage
+import settings
 
 test_app = flask.Flask(__name__)
 
 
 class OriginTrialsClientTest(testing_config.CustomTestCase):
 
-  mock_list_trials_json = {
-    'trials': [
-      {
-        'id': '-5269211564023480319',
-        'displayName': 'Example Trial',
-        'description': 'A description.',
-        'originTrialFeatureName': 'ExampleTrial',
-        'status': 'ACTIVE',
-        'enabled': True,
-        'isPublic': True,
-        'chromestatusUrl': 'https://example.com/chromestatus',
-        'startMilestone': '123',
-        'endMilestone': '456',
-        'originalEndMilestone': '450',
-        'endTime': '2025-01-01T00:00:00Z',
-        'feedbackUrl': 'https://example.com/feedback',
-        'documentationUrl': 'https://example.com/docs',
-        'intentToExperimentUrl': 'https://example.com/intent',
-        'type': 'ORIGIN_TRIAL',
-        'allowThirdPartyOrigins': True,
-        'trialExtensions': [{}],
-      },
-      {
-        'id': '3611886901151137793',
-        'displayName': 'Non-public trial',
-        'description': 'Another description.',
-        'originTrialFeatureName': 'SampleTrial',
-        'status': 'COMPLETE',
-        'enabled': True,
-        'isPublic': False,
-        'chromestatusUrl': 'https://example.com/chromestatus2',
-        'startMilestone': '100',
-        'endMilestone': '200',
-        'endTime': '2024-01-01T00:00:00Z',
-      }
-    ]
-  }
+  def setUp(self):
+    self.ot_stage = Stage(
+        feature_id=1, stage_type=150, ot_display_name='Example Trial',
+        milestones=MilestoneSet(desktop_first=100, desktop_last=106),
+        ot_documentation_url='https://example.com/docs',
+        ot_feedback_submission_url='https://example.com/feedback',
+        intent_thread_url='https://example.com/experiment',
+        ot_description='OT description', ot_has_third_party_support=True,
+        ot_is_deprecation_trial=True)
+    self.ot_stage.put()
+    self.mock_list_trials_json = {
+      'trials': [
+        {
+          'id': '-5269211564023480319',
+          'displayName': 'Example Trial',
+          'description': 'A description.',
+          'originTrialFeatureName': 'ExampleTrial',
+          'status': 'ACTIVE',
+          'enabled': True,
+          'isPublic': True,
+          'chromestatusUrl': 'https://example.com/chromestatus',
+          'startMilestone': '123',
+          'endMilestone': '456',
+          'originalEndMilestone': '450',
+          'endTime': '2025-01-01T00:00:00Z',
+          'feedbackUrl': 'https://example.com/feedback',
+          'documentationUrl': 'https://example.com/docs',
+          'intentToExperimentUrl': 'https://example.com/intent',
+          'type': 'ORIGIN_TRIAL',
+          'allowThirdPartyOrigins': True,
+          'trialExtensions': [{}],
+        },
+        {
+          'id': '3611886901151137793',
+          'displayName': 'Non-public trial',
+          'description': 'Another description.',
+          'originTrialFeatureName': 'SampleTrial',
+          'status': 'COMPLETE',
+          'enabled': True,
+          'isPublic': False,
+          'chromestatusUrl': 'https://example.com/chromestatus2',
+          'startMilestone': '100',
+          'endMilestone': '200',
+          'endTime': '2024-01-01T00:00:00Z',
+        }
+      ]
+    }
+
+  def tearDown(self):
+    for entity in Stage.query():
+      entity.key.delete()
 
   @mock.patch('framework.secrets.get_ot_api_key')
   @mock.patch('requests.get')
@@ -131,7 +147,7 @@ def test_extend_origin_trial__no_api_key(
   def test_extend_origin_trial__with_api_key(
       self, mock_requests_post, mock_get_trial_end_time,
       mock_get_ot_access_token, mock_api_key_get):
-    """If an API key is available, POST should extend trial and return true."""
+    """If an API key is available, POST should extend trial."""
     mock_requests_post.return_value = mock.MagicMock(
         status_code=200, json=lambda : {})
     mock_get_trial_end_time.return_value = 111222333
@@ -159,3 +175,94 @@ def test_get_trial_end_time(self, mock_requests_get):
     return_result = origin_trials_client._get_trial_end_time(123)
     self.assertEqual(return_result, 1682812800)
     mock_requests_get.assert_called_once()
+
+  @mock.patch('framework.secrets.get_ot_api_key')
+  @mock.patch('requests.post')
+  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)
+
+    mock_api_key_get.assert_called_once()
+    # POST request should not be executed with no API key.
+    mock_requests_post.assert_not_called()
+
+  @mock.patch('framework.secrets.get_ot_api_key')
+  @mock.patch('framework.origin_trials_client._get_ot_access_token')
+  @mock.patch('framework.origin_trials_client._get_trial_end_time')
+  @mock.patch('requests.post')
+  def test_create_origin_trial__with_api_key(
+      self, mock_requests_post, mock_get_trial_end_time,
+      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'})
+    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')
+
+    mock_api_key_get.assert_called_once()
+    mock_get_ot_access_token.assert_called_once()
+    mock_requests_post.assert_called_once_with(
+      f'{settings.OT_API_URL}/v1/trials-integration',
+      headers={'Authorization': 'Bearer access_token'},
+      params={'key': 'api_key_value'},
+      json={
+        'trial': {
+          'display_name': 'Example Trial',
+          'start_milestone': '100',
+          'end_milestone': '106',
+          'end_time': {
+            'seconds': 111222333
+          },
+          'description': 'OT description',
+          'documentation_url': 'https://example.com/docs',
+          'feedback_url': 'https://example.com/feedback',
+          'intent_to_experiment_url': 'https://example.com/experiment',
+          'chromestatus_url': f'{settings.SITE_URL}feature/1',
+          'allow_third_party_origins': True,
+          'type': 'DEPRECATION',
+        }
+      }
+    )
+
+  @mock.patch('framework.secrets.get_ot_api_key')
+  @mock.patch('requests.post')
+  def test_activate_origin_trial__no_api_key(
+      self, mock_requests_post, mock_api_key_get):
+    """If no API key is available, do not send activation request."""
+    mock_api_key_get.return_value = None
+    origin_trials_client.create_origin_trial(self.ot_stage)
+
+    mock_api_key_get.assert_called_once()
+    # POST request should not be executed with no API key.
+    mock_requests_post.assert_not_called()
+
+  @mock.patch('framework.secrets.get_ot_api_key')
+  @mock.patch('framework.origin_trials_client._get_ot_access_token')
+  @mock.patch('framework.origin_trials_client._get_trial_end_time')
+  @mock.patch('requests.post')
+  def test_activate_origin_trial__with_api_key(
+      self, mock_requests_post, mock_get_trial_end_time,
+      mock_get_ot_access_token, mock_api_key_get):
+    """If an API key is available, POST should activate trial."""
+    mock_requests_post.return_value = mock.MagicMock(
+        status_code=200, json=lambda : {})
+    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'
+
+    origin_trials_client.activate_origin_trial('-1234567890')
+
+    mock_api_key_get.assert_called_once()
+    mock_get_ot_access_token.assert_called_once()
+    mock_requests_post.assert_called_once_with(
+      f'{settings.OT_API_URL}/v1/trials-integration/-1234567890:activate',
+      headers={'Authorization': 'Bearer access_token'},
+      params={'key': 'api_key_value'},
+      json={'id': '-1234567890'}
+    )

framework/utils.py

@@ -15,14 +15,15 @@
 
 import calendar
 import datetime
-import flask
 import logging
+import requests
 import time
 import traceback
 
-from framework import users
 import settings
 
+CHROMIUM_SCHEDULE_DATE_FORMAT = '%Y-%m-%dT%H:%M:%S'
+
 
 def normalized_name(val):
   return val.lower().replace(' ', '').replace('/', '')
@@ -117,3 +118,15 @@ def get_banner_time(timestamp):
 def dedupe(list_with_duplicates):
   """Return a list without duplicates, in the original order."""
   return list(dict.fromkeys(list_with_duplicates))
+
+
+def get_chromium_milestone_info(milestone: int) -> dict:
+  try:
+    response = requests.get(
+      'https://chromiumdash.appspot.com/fetch_milestone_schedule'
+      f'?mstone={milestone}')
+    response.raise_for_status()
+  except requests.exceptions.RequestException as e:
+    logging.exception('Failed to get response from Chromium schedule API.')
+    raise e
+  return response.json()