feat: enable new TA tasks on new repos (#108)

joseph-sentry · web-flow · commit c047c24cfeff · 2025-06-19T14:08:02.000Z
* feat: enable new TA tasks on new repos

* update the date

* fix tests
diff --git a/apps/worker/services/test_analytics/ta_metrics.py b/apps/worker/services/test_analytics/ta_metrics.py
@@ -1,4 +1,4 @@
-from shared.metrics import Summary
+from shared.metrics import Counter, Summary
 
 write_tests_summary = Summary(
     "write_tests_summary",
@@ -37,3 +37,8 @@
     "The time it takes to process flakes",
     ["impl"],
 )
+
+new_ta_tasks_repo_summary = Counter(
+    "new_ta_tasks_repo_summary",
+    "Number of repos that were processed using 'new' ta tasks, due to being created after 2025-05-30",
+)
diff --git a/apps/worker/tasks/tests/unit/test_upload_task.py b/apps/worker/tasks/tests/unit/test_upload_task.py
@@ -29,7 +29,7 @@
 from tasks.bundle_analysis_processor import bundle_analysis_processor_task
 from tasks.test_results_finisher import test_results_finisher_task
 from tasks.test_results_processor import test_results_processor_task
-from tasks.upload import UploadContext, UploadTask
+from tasks.upload import NEW_TA_TASKS_CUTOFF_DATE, UploadContext, UploadTask
 from tasks.upload_finisher import upload_finisher_task
 from tasks.upload_processor import upload_processor_task
 
@@ -112,6 +112,45 @@ def mock_redis(mocker):
     yield redis_server
 
 
+@pytest.fixture()
+def mock_repo_provider_service(mocker):
+    mocker.patch("tasks.upload.hasattr", return_value=False)
+
+    mock_repo_provider_service = AsyncMock()
+    mock_repo_provider_service.get_commit.return_value = {
+        "author": {
+            "id": "123",
+            "username": "456",
+            "email": "789",
+            "name": "101",
+        },
+        "message": "hello world",
+        "parents": [],
+        "timestamp": str(datetime.now()),
+    }
+    mock_repo_provider_service.get_ancestors_tree.return_value = {"parents": []}
+    mock_repo_provider_service.get_pull_request.return_value = {
+        "head": {"branch": "main"},
+        "base": {},
+    }
+    mock_repo_provider_service.list_top_level_files.return_value = [
+        {"name": "codecov.yml", "path": "codecov.yml"}
+    ]
+    mock_repo_provider_service.get_source.return_value = {
+        "content": """
+        codecov:
+            max_report_age: 1y ago
+        """
+    }
+
+    mocker.patch(
+        "tasks.base.get_repo_provider_service",
+        return_value=mock_repo_provider_service,
+    )
+
+    return mock_repo_provider_service
+
+
 @pytest.fixture(scope="function", autouse=True)
 def clear_log_context(mocker):
     set_log_context(LogContext())
@@ -442,6 +481,7 @@ def test_upload_task_call_new_ta_tasks(
         mock_storage,
         mock_redis,
         celery_app,
+        mock_repo_provider_service,
     ):
         chain = mocker.patch("tasks.upload.chain")
         _ = mocker.patch("tasks.upload.NEW_TA_TASKS.check_value", return_value="both")
@@ -450,38 +490,6 @@ def test_upload_task_call_new_ta_tasks(
         jsonified_redis_queue = [json.dumps(x) for x in redis_queue]
         mocker.patch.object(UploadTask, "app", celery_app)
 
-        mock_repo_provider_service = AsyncMock()
-        mock_repo_provider_service.get_commit.return_value = {
-            "author": {
-                "id": "123",
-                "username": "456",
-                "email": "789",
-                "name": "101",
-            },
-            "message": "hello world",
-            "parents": [],
-            "timestamp": str(datetime.now()),
-        }
-        mock_repo_provider_service.get_ancestors_tree.return_value = {"parents": []}
-        mock_repo_provider_service.get_pull_request.return_value = {
-            "head": {"branch": "main"},
-            "base": {},
-        }
-        mock_repo_provider_service.list_top_level_files.return_value = [
-            {"name": "codecov.yml", "path": "codecov.yml"}
-        ]
-        mock_repo_provider_service.get_source.return_value = {
-            "content": """
-            codecov:
-                max_report_age: 1y ago
-            """
-        }
-
-        mocker.patch(
-            "tasks.base.get_repo_provider_service",
-            return_value=mock_repo_provider_service,
-        )
-        mocker.patch("tasks.upload.hasattr", return_value=False)
         commit = CommitFactory.create(
             message="",
             commitid="abf6d4df662c47e32460020ab14abf9303581429",
@@ -542,6 +550,91 @@ def test_upload_task_call_new_ta_tasks(
         notify_sig = test_results_finisher_task.signature(kwargs=kwargs)
         chain.assert_has_calls([call(processor_sig, notify_sig)], any_order=True)
 
+    @pytest.mark.django_db(databases={"default"}, transaction=True)
+    def test_upload_task_call_test_results_new_repo(
+        self,
+        mocker,
+        mock_configuration,
+        dbsession,
+        codecov_vcr,
+        mock_storage,
+        mock_redis,
+        celery_app,
+        mock_repo_provider_service,
+    ):
+        chain = mocker.patch("tasks.upload.chain")
+        storage_path = "v4/raw/2019-05-22/C3C4715CA57C910D11D5EB899FC86A7E/4c4e4654ac25037ae869caeb3619d485970b6304/a84d445c-9c1e-434f-8275-f18f1f320f81.txt"
+        redis_queue = [{"url": storage_path, "build_code": "some_random_build"}]
+        jsonified_redis_queue = [json.dumps(x) for x in redis_queue]
+
+        commit = CommitFactory.create(
+            message="",
+            commitid="abf6d4df662c47e32460020ab14abf9303581429",
+            repository__owner__oauth_token="GHTZB+Mi+kbl/ubudnSKTJYb/fgN4hRJVJYSIErtidEsCLDJBb8DZzkbXqLujHAnv28aKShXddE/OffwRuwKug==",
+            repository__owner__username="ThiagoCodecov",
+            repository__owner__service="github",
+            repository__yaml={"codecov": {"max_report_age": "1y ago"}},
+            repository__name="example-python",
+            pullid=1,
+            # Setting the time to _before_ patch centric default YAMLs start date of 2024-04-30
+            repository__owner__createstamp=datetime(2023, 1, 1, tzinfo=UTC),
+        )
+        dbsession.add(commit)
+        dbsession.flush()
+        dbsession.refresh(commit)
+        commit.timestamp = NEW_TA_TASKS_CUTOFF_DATE.replace(
+            day=21
+        )  # One day after cutoff
+        dbsession.flush()
+
+        mock_redis.lists[f"uploads/{commit.repoid}/{commit.commitid}/test_results"] = (
+            jsonified_redis_queue
+        )
+
+        mock_repo_provider_service.get_commit.return_value["timestamp"] = (
+            NEW_TA_TASKS_CUTOFF_DATE.replace(day=21)
+        )  # One day after cutoff
+
+        UploadTask().run_impl(
+            dbsession,
+            commit.repoid,
+            commit.commitid,
+            report_type="test_results",
+        )
+
+        commit_report = commit.commit_report(report_type=ReportType.TEST_RESULTS)
+        assert commit_report
+        uploads = commit_report.uploads
+        assert len(uploads) == 1
+        upload = dbsession.query(Upload).filter_by(report_id=commit_report.id).first()
+        processor_sig = test_results_processor_task.s(
+            False,
+            repoid=commit.repoid,
+            commitid=commit.commitid,
+            commit_yaml={"codecov": {"max_report_age": "1y ago"}},
+            arguments_list=[
+                {
+                    "url": storage_path,
+                    "flags": [],
+                    "build_code": "some_random_build",
+                    "upload_id": upload.id,
+                    "upload_pk": upload.id,
+                }
+            ],
+            impl_type="new",
+        )
+        kwargs = {
+            "repoid": commit.repoid,
+            "commitid": commit.commitid,
+            "commit_yaml": {"codecov": {"max_report_age": "1y ago"}},
+            "checkpoints_TestResultsFlow": None,
+            "impl_type": "new",
+        }
+
+        kwargs[_kwargs_key(TestResultsFlow)] = mocker.ANY
+        notify_sig = test_results_finisher_task.signature(kwargs=kwargs)
+        chain.assert_called_with(*[processor_sig, notify_sig])
+
     def test_upload_task_call_no_jobs(
         self,
         mocker,
diff --git a/apps/worker/tasks/upload.py b/apps/worker/tasks/upload.py
@@ -3,14 +3,14 @@
 import time
 import uuid
 from copy import deepcopy
+from datetime import UTC, datetime
 from typing import TypedDict
 
 import orjson
 import sentry_sdk
 from asgiref.sync import async_to_sync
 from celery import chain, chord
 from django.conf import settings
-from django.utils import timezone
 from redis import Redis
 from redis.exceptions import LockError
 from sqlalchemy.orm import Session
@@ -36,6 +36,7 @@
     gitlab_webhook_update,
     possibly_update_commit_from_provider_info,
 )
+from services.test_analytics.ta_metrics import new_ta_tasks_repo_summary
 from services.test_results import TestResultsReportService
 from shared.celery_config import upload_task_name
 from shared.config import get_config
@@ -56,6 +57,9 @@
 
 log = logging.getLogger(__name__)
 
+# Date after which commits are considered for new TA tasks implementation
+NEW_TA_TASKS_CUTOFF_DATE = datetime(2025, 6, 20, tzinfo=UTC)
+
 CHUNK_SIZE = 3
 
 UPLOADS_PER_TASK_SCHEDULE = Histogram(
@@ -539,7 +543,7 @@ def _possibly_create_uploads_to_insert(
 
         # List to track possible measurements to insert later
         measurements_list: list[UserMeasurement] = []
-        created_at = timezone.now()
+        created_at = datetime.now(UTC)
 
         # List + helper mapping to track possible upload + flags to insert later
         upload_flag_map: dict[Upload, list | str | None] = {}
@@ -750,6 +754,22 @@ def _schedule_ta_processing_task(
         new_ta_tasks = NEW_TA_TASKS.check_value(commit.repoid, default="old")
         if not settings.TA_TIMESERIES_ENABLED:
             new_ta_tasks = "old"
+        else:
+            db_session: Session = commit.get_db_session()  # type: ignore
+            earliest_commit = (
+                db_session.query(Commit)
+                .filter(
+                    Commit.repoid == commit.repoid,
+                    Commit.timestamp > NEW_TA_TASKS_CUTOFF_DATE,
+                )
+                .order_by(Commit.timestamp)
+                .limit(1)
+                .first()
+            )
+
+            if earliest_commit:
+                new_ta_tasks = "new"
+                new_ta_tasks_repo_summary.inc()
 
         task_group = [
             test_results_processor_task.s(
