From 5fed1ade14d83fb036ccb727f9129a0feee557e4 Mon Sep 17 00:00:00 2001
From: Daniel Smith <56164590+DanielRyanSmith@users.noreply.github.com>
Date: Thu, 16 Jan 2025 21:08:10 +0000
Subject: [PATCH] New histogram ID enum (#4703)

Co-authored-by: DanielRyanSmith <danielrsmith@google.com>
---
 framework/origin_trials_client.py      | 13 ++++++++-----
 framework/origin_trials_client_test.py |  7 +++----
 internals/core_enums.py                |  7 +++++++
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/framework/origin_trials_client.py b/framework/origin_trials_client.py
index b1c427cfdfcb..65d4c3ca0ee5 100644
--- a/framework/origin_trials_client.py
+++ b/framework/origin_trials_client.py
@@ -23,6 +23,7 @@
 
 from framework import secrets
 from framework import utils
+from internals.core_enums import BlinkHistogramID
 from internals.core_models import Stage
 from internals.data_types import OriginTrialInfo
 import settings
@@ -30,6 +31,7 @@
 
 class UseCounterConfig(TypedDict):
   bucket_number: int
+  histogram_id: str
 
 class RequestTrial(TypedDict):
   id: NotRequired[int]
@@ -46,7 +48,6 @@ class RequestTrial(TypedDict):
   type: str
   origin_trial_feature_name: NotRequired[str]
   blink_use_counter_config: NotRequired[UseCounterConfig]
-  blink_webdx_use_counter_config: NotRequired[UseCounterConfig]
 
 
 class InternalRegistrationConfig(TypedDict):
@@ -184,12 +185,14 @@ def _send_create_trial_request(
   if ot_stage.ot_is_deprecation_trial:
     json['registration_config']['allow_public_suffix_subdomains'] = True
   if ot_stage.ot_use_counter_bucket_number:
-    config: UseCounterConfig = {'bucket_number': ot_stage.ot_use_counter_bucket_number}
+    config: UseCounterConfig = {
+      'bucket_number': ot_stage.ot_use_counter_bucket_number,
+      'histogram_id': BlinkHistogramID.web_feature.value
+    }
     if (ot_stage.ot_chromium_trial_name
         and ot_stage.ot_chromium_trial_name.startswith('WebDXFeature::')):
-      json['trial']['blink_webdx_use_counter_config'] = config
-    else:
-      json['trial']['blink_use_counter_config'] = config
+      config['histogram_id'] = BlinkHistogramID.webdx_feature.value
+    json['trial']['blink_use_counter_config'] = config
 
   headers = {'Authorization': f'Bearer {access_token}'}
   url = f'{settings.OT_API_URL}/v1/trials:initialize'
diff --git a/framework/origin_trials_client_test.py b/framework/origin_trials_client_test.py
index 69a05feb6359..ef7732f7d20c 100644
--- a/framework/origin_trials_client_test.py
+++ b/framework/origin_trials_client_test.py
@@ -240,6 +240,7 @@ def test_create_origin_trial__with_api_key(
           'type': 'DEPRECATION',
           'blink_use_counter_config': {
             'bucket_number': 11,
+            'histogram_id': 'WEB_FEATURE',
           }
         }, create_trial_json['trial'])
     self.assertEqual({
@@ -287,11 +288,9 @@ def test_create_origin_trial__webdx_feature(
     # 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']
-    # WebFeature config should be null.
-    self.assertEqual(None, create_trial_json['trial'].get('blink_use_counter_config'))
     # WebDXFeature config should be populated.
-    self.assertEqual({'bucket_number': 11},
-                     create_trial_json['trial']['blink_webdx_use_counter_config'])
+    self.assertEqual({'bucket_number': 11, 'histogram_id': 'WEBDX_FEATURE'},
+                     create_trial_json['trial']['blink_use_counter_config'])
 
   @mock.patch('framework.secrets.get_ot_api_key')
   @mock.patch('requests.post')
diff --git a/internals/core_enums.py b/internals/core_enums.py
index 83910e33a432..a0d7848504d5 100644
--- a/internals/core_enums.py
+++ b/internals/core_enums.py
@@ -15,6 +15,7 @@
 
 import collections
 import re
+from enum import Enum
 from typing import Optional
 
 
@@ -503,6 +504,12 @@
 OT_CREATED = 5
 OT_ACTIVATED = 6
 
+# Histogram IDs used for identifying origin trial use counters.
+class BlinkHistogramID(str, Enum):
+  web_feature = 'WEB_FEATURE'
+  webdx_feature = 'WEBDX_FEATURE'
+  css_property_id = 'CSS_PROPERTY_ID'
+
 NO_ACTIVE_DEV = 1
 PROPOSED = 2
 IN_DEVELOPMENT = 3