Support production/testing mozci environments in ClassificationLoader regex (#7419)

Co-authored-by: Eva Bardou <ebardou@teklia.com>

tests/etl/test_classification_loader.py

@@ -126,22 +126,45 @@ def populate_bugscache():
 
 
 @pytest.mark.parametrize(
-    'route',
+    'mode, route',
     [
-        'completely bad route',
-        'index.project.mozci.classification..revision.A35mWTRuQmyj88yMnIF0fA',
-        'index.project.mozci.classification.autoland.revision.',
-        'index.project.mozci.classification.autoland.revision.-35mW@RuQ__j88yénIF0f-',
+        ('production', 'completely bad route'),
+        ('production', 'index.project.mozci.classification..revision.A35mWTRuQmyj88yMnIF0fA'),
+        ('production', 'index.project.mozci.classification.autoland.revision.'),
+        (
+            'production',
+            'index.project.mozci.classification.autoland.revision.-35mW@RuQ__j88yénIF0f-',
+        ),
+        (
+            'production',
+            'index.project.mozci.testing.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA',
+        ),
+        ('testing', 'index.project.mozci.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA'),
     ],
 )
-def test_get_push_wrong_route(route):
+def test_get_push_wrong_route(mode, route, monkeypatch):
+    monkeypatch.setenv('PULSE_MOZCI_ENVIRONMENT', mode)
+
     with pytest.raises(AttributeError):
         ClassificationLoader().get_push(route)
 
 
 @pytest.mark.django_db
-def test_get_push_unsupported_project():
-    route = 'index.project.mozci.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA'
+@pytest.mark.parametrize(
+    'mode, route',
+    [
+        (
+            'production',
+            'index.project.mozci.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA',
+        ),
+        (
+            'testing',
+            'index.project.mozci.testing.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA',
+        ),
+    ],
+)
+def test_get_push_unsupported_project(mode, route, monkeypatch):
+    monkeypatch.setenv('PULSE_MOZCI_ENVIRONMENT', mode)
 
     with pytest.raises(Repository.DoesNotExist) as e:
         ClassificationLoader().get_push(route)
@@ -150,8 +173,21 @@ def test_get_push_unsupported_project():
 
 
 @pytest.mark.django_db
-def test_get_push_unsupported_revision(autoland_repository):
-    route = 'index.project.mozci.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA'
+@pytest.mark.parametrize(
+    'mode, route',
+    [
+        (
+            'production',
+            'index.project.mozci.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA',
+        ),
+        (
+            'testing',
+            'index.project.mozci.testing.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA',
+        ),
+    ],
+)
+def test_get_push_unsupported_revision(mode, route, autoland_repository, monkeypatch):
+    monkeypatch.setenv('PULSE_MOZCI_ENVIRONMENT', mode)
 
     with pytest.raises(Push.DoesNotExist) as e:
         ClassificationLoader().get_push(route)
@@ -160,8 +196,21 @@ def test_get_push_unsupported_revision(autoland_repository):
 
 
 @pytest.mark.django_db
-def test_get_push(autoland_push):
-    route = 'index.project.mozci.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA'
+@pytest.mark.parametrize(
+    'mode, route',
+    [
+        (
+            'production',
+            'index.project.mozci.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA',
+        ),
+        (
+            'testing',
+            'index.project.mozci.testing.classification.autoland.revision.A35mWTRuQmyj88yMnIF0fA',
+        ),
+    ],
+)
+def test_get_push(mode, route, autoland_push, monkeypatch):
+    monkeypatch.setenv('PULSE_MOZCI_ENVIRONMENT', mode)
 
     assert ClassificationLoader().get_push(route) == autoland_push
 

treeherder/etl/classification_loader.py

@@ -19,9 +19,12 @@ from treeherder.utils.taskcluster import download_artifact, get_task_definition
 env = environ.Env()
 logger = logging.getLogger(__name__)
 
-CLASSIFICATION_ROUTE_REGEX = re.compile(
+CLASSIFICATION_PRODUCTION_ROUTE_REGEX = re.compile(
     r"index\.project\.mozci\.classification\.(.+)\.revision\.([0-9A-Fa-f]+)"
 )
+CLASSIFICATION_TESTING_ROUTE_REGEX = re.compile(
+    r"index\.project\.mozci\.testing\.classification\.(.+)\.revision\.([0-9A-Fa-f]+)"
+)
 
 
 class ClassificationLoader:
@@ -82,8 +85,18 @@ class ClassificationLoader:
         )
 
     def get_push(self, task_route):
+        mozci_env = env('PULSE_MOZCI_ENVIRONMENT', default='production')
+        if mozci_env == 'testing':
+            route_regex = CLASSIFICATION_TESTING_ROUTE_REGEX
+        else:
+            if mozci_env != 'production':
+                logger.warning(
+                    f'PULSE_MOZCI_ENVIRONMENT should be testing or production not {mozci_env}, defaulting to production'
+                )
+            route_regex = CLASSIFICATION_PRODUCTION_ROUTE_REGEX
+
         try:
-            project, revision = CLASSIFICATION_ROUTE_REGEX.search(task_route).groups()
+            project, revision = route_regex.search(task_route).groups()
         except AttributeError as e:
             logger.error(
                 "Failed to parse the given route '%s' to retrieve the push project and revision: %s",