Disable OT milestone edits during automated OT creation process (server-side) (#4236)

* No OT milestone edits during OT creation process

* fix imports

api/stages_api_test.py

@@ -22,9 +22,10 @@ from google.cloud import ndb  # type: ignore
 import werkzeug.exceptions
 
 from api import stages_api
-from internals.user_models import AppUser
+from internals.core_enums import OT_READY_FOR_CREATION
 from internals.core_models import FeatureEntry, MilestoneSet, Stage
 from internals.review_models import Gate
+from internals.user_models import AppUser
 
 test_app = flask.Flask(__name__)
 
@@ -565,6 +566,53 @@ class StagesAPITest(testing_config.CustomTestCase):
 
     self.assertEqual(actual, 'fake response')
 
+  @mock.patch('flask.abort')
+  def test_patch__ot_milestones_during_creation(self, mock_abort):
+    """Raises 400 if OT start milestone is updated during OT creation process.
+    """
+    testing_config.sign_in('feature_owner@example.com', 123)
+    json = {
+      'desktop_first': {
+        'form_field_name': 'ot_milestone_desktop_start',
+        'value': 200,
+      },
+    }
+
+    # OT is flagged for automated creation process.
+    self.stage_1.ot_setup_status = OT_READY_FOR_CREATION
+    self.stage_1.put()
+    mock_abort.side_effect = werkzeug.exceptions.BadRequest
+    with test_app.test_request_context(
+        f'{self.request_path}1/stages/10', json=json):
+      with self.assertRaises(werkzeug.exceptions.BadRequest):
+        self.handler.do_patch(feature_id=1, stage_id=10)
+    mock_abort.assert_called_once_with(
+        400,
+        description='Cannot edit OT milestones while creation is in progress.')
+
+  @mock.patch('flask.abort')
+  def test_patch__ot_end_milestone_during_creation(self, mock_abort):
+    """Raises 400 if OT end milestone is updated during OT creation process."""
+    testing_config.sign_in('feature_owner@example.com', 123)
+    json = {
+      'desktop_last': {
+        'form_field_name': 'ot_milestone_desktop_end',
+        'value': 206,
+      },
+    }
+
+    # OT is flagged for automated creation process.
+    self.stage_1.ot_setup_status = OT_READY_FOR_CREATION
+    self.stage_1.put()
+    mock_abort.side_effect = werkzeug.exceptions.BadRequest
+    with test_app.test_request_context(
+        f'{self.request_path}1/stages/10', json=json):
+      with self.assertRaises(werkzeug.exceptions.BadRequest):
+        self.handler.do_patch(feature_id=1, stage_id=10)
+    mock_abort.assert_called_once_with(
+        400,
+        description='Cannot edit OT milestones while creation is in progress.')
+
   def test_patch__valid(self):
     """A valid PATCH request should update an existing stage."""
     testing_config.sign_in('feature_owner@example.com', 123)

framework/basehandlers.py

@@ -37,10 +37,14 @@ from framework import utils
 from framework import xsrf
 from internals import approval_defs
 from internals import notifier_helpers
-from internals.core_enums import ALL_ORIGIN_TRIAL_STAGE_TYPES
+from internals import user_models
+from internals.core_enums import (
+  ALL_ORIGIN_TRIAL_STAGE_TYPES,
+  OT_ACTIVATION_FAILED,
+  OT_CREATION_FAILED,
+  OT_READY_FOR_CREATION)
 from internals.core_models import FeatureEntry, MilestoneSet, Stage
 from internals.data_types import CHANGED_FIELDS_LIST_TYPE
-from internals import user_models
 
 from flask import session
 from flask import render_template
@@ -354,7 +358,18 @@ class EntitiesAPIHandler(APIHandler):
     """Update stage fields with changes provided."""
     stage_was_updated = False
     ot_action_requested = False
-    # Check if valid ID is provided and fetch stage if it exists.
+
+    mutating_ot_milestones = any(
+        v['form_field_name'] == 'ot_milestone_desktop_start' or
+        v['form_field_name'] == 'ot_milestone_desktop_end'
+        for v in change_info.values())
+    ot_creation_in_progress =  (
+        stage.ot_setup_status == OT_READY_FOR_CREATION or
+        stage.ot_setup_status == OT_CREATION_FAILED or
+        stage.ot_setup_status == OT_ACTIVATION_FAILED)
+    if mutating_ot_milestones and ot_creation_in_progress:
+      self.abort(400,
+                 'Cannot edit OT milestones while creation is in progress.')
 
     # Update stage fields.
     for field, field_type in api_specs.STAGE_FIELD_DATA_TYPES: