diff --git a/internals/maintenance_scripts.py b/internals/maintenance_scripts.py index 7e6f9eb451ee..2651d98ce575 100644 --- a/internals/maintenance_scripts.py +++ b/internals/maintenance_scripts.py @@ -230,52 +230,93 @@ def get_template_data(self, **kwargs): class AssociateOTs(FlaskHandler): - def write_fields_for_trial_stage(self, trial_stage: Stage, trial_data: dict[str, Any]): - """Check if any OT stage fields are unfilled and populate them with - the matching trial data. + def write_field( + self, + trial_stage: Stage, + trial_data: dict[str, Any], + stage_field_name: str, + trial_field_name: str, + ) -> bool: + """Set the OT stage value to the value from the OT console if it is unset. + + Returns: + boolean value of whether or not the value was changed on the stage. """ - if trial_stage.origin_trial_id is None: - trial_stage.origin_trial_id = trial_data['id'] - - if trial_stage.ot_chromium_trial_name is None: - trial_stage.ot_chromium_trial_name = trial_data['origin_trial_feature_name'] + if (not getattr(trial_stage, stage_field_name) + and trial_data[trial_field_name]): + setattr(trial_stage, stage_field_name, trial_data[trial_field_name]) + return True + return False + def write_milestone_field( + self, + trial_stage: Stage, + trial_data: dict[str, Any], + stage_field_name: str, + trial_field_name: str, + ) -> bool: + """Set an OT milestone value to the value from the OT console + if it is unset. + + Returns: + boolean value of whether or not the value was changed on the stage. + """ if trial_stage.milestones is None: trial_stage.milestones = MilestoneSet() - if (trial_stage.milestones.desktop_first is None and - trial_data['start_milestone'] is not None): - trial_stage.milestones.desktop_first = int(trial_data['start_milestone']) - if trial_stage.milestones.desktop_last is None: - # An original end milestone is kept if the trial has had extensions. - # TODO(DanielRyanSmith): Extension milestones in the trial data - # should be associated with new extension stages for data accuracy. - if trial_data['original_end_milestone'] is not None: - trial_stage.milestones.desktop_last = ( - int(trial_data['original_end_milestone'])) - elif trial_data['end_milestone'] is not None: - trial_stage.milestones.desktop_last = ( - int(trial_data['end_milestone'])) - - if trial_stage.display_name is None: - trial_stage.display_name = trial_data['display_name'] - - if trial_stage.intent_thread_url is None: - trial_stage.intent_thread_url = trial_data['intent_to_experiment_url'] - - if trial_stage.ot_feedback_submission_url is None: - trial_stage.ot_feedback_submission_url = trial_data['feedback_url'] - - if trial_stage.ot_documentation_url is None: - trial_stage.ot_documentation_url = trial_data['documentation_url'] - - if trial_stage.ot_has_third_party_support: - trial_stage.ot_has_third_party_support = trial_data['allow_third_party_origins'] + if (getattr(trial_stage.milestones, stage_field_name) is None and + trial_data[trial_field_name] is not None): + setattr(trial_stage.milestones, + stage_field_name, int(trial_data[trial_field_name])) + return True + return False + + def write_fields_for_trial_stage( + self, trial_stage: Stage, trial_data: dict[str, Any]): + """Check if any OT stage fields are unfilled and populate them with + the matching trial data. + """ + stage_changed = False + stage_changed = (self.write_field( + trial_stage, trial_data, 'origin_trial_id', 'id') or stage_changed) + stage_changed = (self.write_field( + trial_stage, trial_data, + 'ot_chromium_trial_name', 'origin_trial_feature_name') or stage_changed) + stage_changed = (self.write_field( + trial_stage, trial_data, + 'ot_feedback_submission_url', 'feedback_url') or stage_changed) + stage_changed = (self.write_field( + trial_stage, trial_data, + 'ot_documentation_url', 'documentation_url') or stage_changed) + stage_changed = (self.write_field( + trial_stage, trial_data, + 'intent_thread_url', 'intent_to_experiment_url') or stage_changed) + stage_changed = (self.write_field( + trial_stage, trial_data, + 'display_name', 'display_name') or stage_changed) + stage_changed = (self.write_field( + trial_stage, trial_data, + 'ot_display_name', 'display_name') or stage_changed) + stage_changed = (self.write_field( + trial_stage, trial_data, + 'ot_has_third_party_support', 'allow_third_party_origins') or + stage_changed) + stage_changed = (self.write_milestone_field( + trial_stage, trial_data, + 'desktop_first', 'start_milestone') or stage_changed) + stage_changed = (self.write_milestone_field( + trial_stage, trial_data, + 'desktop_last', 'original_end_milestone') or stage_changed) if not trial_stage.ot_is_deprecation_trial: trial_stage.ot_is_deprecation_trial = trial_data['type'] == 'DEPRECATION' + stage_changed = True # Clear the trial creation request if it's active. - trial_stage.ot_action_requested = False + if trial_stage.ot_action_requested: + trial_stage.ot_action_requested = False + stage_changed = True + + return stage_changed def parse_feature_id(self, chromestatus_url: str|None) -> int|None: if chromestatus_url is None: @@ -320,12 +361,34 @@ def find_trial_stage(self, feature_id: int) -> Stage|None: return None return trial_stages[0] + def clear_extension_requests(self, ot_stage: Stage, trial_data: dict) -> int: + """Clear any trial extension requests if they have been processed""" + extension_stages: list[Stage] = Stage.query( + Stage.ot_action_requested == True, + Stage.ot_stage_id == ot_stage.key.integer_id()).fetch() + if len(extension_stages) == 0: + return 0 + extension_stages_to_update = [] + for extension_stage in extension_stages: + extension_end = extension_stage.milestones.desktop_last + # If the end milestone of the trial is equal or greater than the + # requested end milestone on the extension stage, we can assume the + # extension request has been processed. + if (int(trial_data['end_milestone']) >= extension_end): + extension_stage.ot_action_requested = False + extension_stages_to_update.append(extension_stage) + + if extension_stages_to_update: + ndb.put_multi(extension_stages_to_update) + return len(extension_stages_to_update) + def get_template_data(self, **kwargs): """Link existing origin trials with their ChromeStatus entry""" self.require_cron_header() trials_list = origin_trials_client.get_trials_list() entities_to_write: list[Stage] = [] + extensions_cleared = 0 # Keep track of stages we're writing to so we avoid trying to write # to the same Stage entity twice in the same batch. unique_entities_to_write: set[int] = set() @@ -339,11 +402,14 @@ def get_template_data(self, **kwargs): continue # If this trial is already associated with a ChromeStatus stage, - # just see if any unfilled fields need to be populated. + # just see if any unfilled fields need to be populated and clear + # any pending extension requests. if stage: - self.write_fields_for_trial_stage(stage, trial_data) - unique_entities_to_write.add(stage.key.integer_id()) - entities_to_write.append(stage) + stage_changed = self.write_fields_for_trial_stage(stage, trial_data) + if stage_changed: + unique_entities_to_write.add(stage.key.integer_id()) + entities_to_write.append(stage) + extensions_cleared += self.clear_extension_requests(stage, trial_data) continue feature_id = self.parse_feature_id(trial_data['chromestatus_url']) @@ -361,9 +427,10 @@ def get_template_data(self, **kwargs): logging.info(f'Already writing to Stage entity {ot_stage_id}') continue - self.write_fields_for_trial_stage(ot_stage, trial_data) - unique_entities_to_write.add(ot_stage_id) - entities_to_write.append(ot_stage) + stage_changed = self.write_fields_for_trial_stage(ot_stage, trial_data) + if stage_changed: + unique_entities_to_write.add(ot_stage_id) + entities_to_write.append(ot_stage) # List any origin trials that did not get associated with a feature entry. if len(trials_with_no_feature) > 0: @@ -373,11 +440,10 @@ def get_template_data(self, **kwargs): for trial_data in trials_with_no_feature: logging.info(f'{trial_data["id"]} {trial_data["display_name"]}') - # Update all the stages at the end. Note that there is a chance - # the stage entities have not changed any values if all fields already - # had a value. + # Update all the stages at the end. logging.info(f'{len(entities_to_write)} stages to update.') if len(entities_to_write) > 0: ndb.put_multi(entities_to_write) - return f'{len(entities_to_write)} Stages updated with trial data.' + return (f'{len(entities_to_write)} Stages updated with trial data.\n' + f'{extensions_cleared} extension requests cleared.') diff --git a/internals/maintenance_scripts_test.py b/internals/maintenance_scripts_test.py new file mode 100644 index 000000000000..3249dacbdc95 --- /dev/null +++ b/internals/maintenance_scripts_test.py @@ -0,0 +1,131 @@ +# Copyright 2023 Google Inc. +# +# 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. +from unittest import mock + +from internals import maintenance_scripts +import testing_config # Must be imported before the module under test. +from internals.core_models import FeatureEntry, Stage, MilestoneSet + +class AssociateOTsTest(testing_config.CustomTestCase): + + def setUp(self): + self.feature_1 = FeatureEntry( + id=123, + name='feature a', summary='sum', category=1, + owner_emails=['feature_owner@example.com']) + self.feature_1.put() + self.feature_2 = FeatureEntry( + id=456, + name='feature b', summary='sum', category=1, + owner_emails=['feature_owner@example.com']) + self.feature_2.put() + + self.ot_stage_1 = Stage( + id=321, feature_id=123, stage_type=150) + self.ot_stage_1.put() + + self.ot_stage_2 = Stage( + id=654, feature_id=456, stage_type=150, origin_trial_id='1') + self.ot_stage_2.put() + + self.extension_stage_1 = Stage( + feature_id=456, + stage_type=151, + ot_stage_id=654, + ot_action_requested=True, + milestones=MilestoneSet(desktop_last=126)) + self.extension_stage_1.put() + + self.trials_list_return_value = [{ + 'id': '4199606652522987521', + 'display_name': 'Sample trial', + 'description': 'Another origin trial', + 'origin_trial_feature_name': 'ChromiumTrialName', + 'enabled': True, + 'status': 'ACTIVE', + 'chromestatus_url': 'https://chromestatus.com/feature/123', + 'start_milestone': '97', + 'end_milestone': '110', + 'original_end_milestone': '100', + 'end_time': '2020-11-10T23:59:59Z', + 'documentation_url': 'https://example.com/docs', + 'feedback_url': 'https://example.com/feedback', + 'intent_to_experiment_url': 'https://example.com/experiment', + 'trial_extensions': [ + { + 'endMilestone': '110', + 'endTime': '2020-11-10T23:59:59Z', + 'extensionIntentUrl': 'https://example.com/extension' + } + ], + 'type': 'DEPRECATION', + 'allow_third_party_origins': True + }, + { + 'id': '1', + 'display_name': 'Sample trial 2', + 'description': 'Another origin trial 2', + 'origin_trial_feature_name': 'ChromiumTrialName2', + 'enabled': True, + 'status': 'ACTIVE', + 'chromestatus_url': 'https://chromestatus.com/feature/456', + 'start_milestone': '120', + 'end_milestone': '126', + 'original_end_milestone': '123', + 'end_time': '2020-11-10T23:59:59Z', + 'documentation_url': 'https://example.com/docs2', + 'feedback_url': 'https://example.com/feedback2', + 'intent_to_experiment_url': 'https://example.com/experiment2', + 'trial_extensions': [ + { + 'endMilestone': '126', + 'endTime': '2020-11-10T23:59:59Z', + 'extensionIntentUrl': 'https://example.com/extension2' + } + ], + 'type': 'ORIGIN_TRIAL', + 'allow_third_party_origins': False + }] + + testing_config.sign_in('one@example.com', 123567890) + + def tearDown(self): + for kind in [FeatureEntry, Stage]: + for entity in kind.query(): + entity.key.delete() + testing_config.sign_out() + + @mock.patch('framework.origin_trials_client.get_trials_list') + def test_associate_ots(self, mock_ot_client): + mock_ot_client.return_value = self.trials_list_return_value + + handler = maintenance_scripts.AssociateOTs() + msg = handler.get_template_data() + self.assertEqual( + msg, '2 Stages updated with trial data.\n1 extension requests cleared.') + + # Check that fields were updated. + ot_1 = self.ot_stage_1 + self.assertEqual(ot_1.ot_display_name, 'Sample trial') + self.assertEqual(ot_1.origin_trial_id, '4199606652522987521') + self.assertEqual(ot_1.ot_chromium_trial_name, 'ChromiumTrialName') + self.assertEqual(ot_1.ot_documentation_url, 'https://example.com/docs') + self.assertEqual(ot_1.intent_thread_url, 'https://example.com/experiment') + self.assertTrue(ot_1.ot_has_third_party_support) + self.assertTrue(ot_1.ot_is_deprecation_trial) + self.assertEqual(ot_1.milestones.desktop_first, 97) + self.assertEqual(ot_1.milestones.desktop_last, 100) + + # Check that the extension request was cleared. + self.assertFalse(self.extension_stage_1.ot_action_requested)