From 54cf4df2e1fce571a6b167c86f21f4726fa04e79 Mon Sep 17 00:00:00 2001
From: Roque Betiol Jr <jbetiol@sig-gis.com>
Date: Fri, 14 Feb 2025 13:57:08 -0300
Subject: [PATCH] Tx Plan execution time metric pushed to Sentry

---
 src/planscape/impacts/tasks.py            | 22 +++++++++++++++
 src/planscape/impacts/tests/test_tasks.py | 34 +++++++++++++++++++++--
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/src/planscape/impacts/tasks.py b/src/planscape/impacts/tasks.py
index ee2d344ab..8b6c01005 100644
--- a/src/planscape/impacts/tasks.py
+++ b/src/planscape/impacts/tasks.py
@@ -1,4 +1,5 @@
 import logging
+from sentry_sdk import metrics
 from typing import Tuple
 from urllib.parse import urljoin
 from celery import chord, chain
@@ -132,6 +133,27 @@ def async_set_status(
         treatment_plan.save()
         log.info(f"Treatment plan {treatment_plan_pk} changed status to {status}.")
 
+    if not start and treatment_plan.started_at and treatment_plan.finished_at:
+        try:
+            elapsed_time_seconds = (
+                treatment_plan.finished_at - treatment_plan.started_at
+            ).total_seconds()
+            log.info(
+                f"Elapsed time for treatment plan {treatment_plan_pk}: {elapsed_time_seconds} seconds."
+            )
+            metrics.gauge(
+                key="tx_plan_elapsed_time",
+                value=elapsed_time_seconds,
+                unit="seconds",
+                tags={
+                    "stand_size": treatment_plan.scenario.get_stand_size(),
+                },
+            )
+        except Exception as e:
+            log.exception(
+                "Failed to send elapsed time metric to Sentry.", extra={"exception": e}
+            )
+
     return (True, treatment_plan_pk)
 
 
diff --git a/src/planscape/impacts/tests/test_tasks.py b/src/planscape/impacts/tests/test_tasks.py
index 5c01882ab..bba08a62b 100644
--- a/src/planscape/impacts/tests/test_tasks.py
+++ b/src/planscape/impacts/tests/test_tasks.py
@@ -9,19 +9,18 @@
 from impacts.tasks import (
     async_send_email_process_finished,
     async_calculate_impacts_for_variable_action_year,
+    async_set_status,
 )
 from stands.models import Stand
 from impacts.services import (
     get_calculation_matrix,
 )
 from impacts.models import (
-    AVAILABLE_YEARS,
     ProjectAreaTreatmentResult,
     TreatmentPrescriptionAction,
     TreatmentResult,
-    ImpactVariable,
+    TreatmentPlanStatus,
 )
-from stands.models import StandMetric
 from datasets.models import DataLayerType
 from datasets.tests.factories import DataLayerFactory
 from impacts.tests.factories import (
@@ -189,3 +188,32 @@ def setUp(self):
         self.project_area = ProjectAreaFactory.create(
             scenario=self.plan.scenario, geometry=self.project_area_geometry
         )
+
+
+class AsyncTestStatusTest(TransactionTestCase):
+    def setUp(self):
+        self.plan = TreatmentPlanFactory.create()
+
+    @mock.patch("impacts.tasks.metrics", autospec=True, return_value=True)
+    def test_status(self, metrics_mock):
+        async_set_status(
+            treatment_plan_pk=self.plan.pk,
+            status=TreatmentPlanStatus.RUNNING,
+            start=True,
+        )
+        self.plan.refresh_from_db()
+        self.assertEquals(self.plan.status, TreatmentPlanStatus.RUNNING)
+        self.assertIsNotNone(self.plan.started_at)
+        self.assertIsNone(self.plan.finished_at)
+        self.assertFalse(metrics_mock.gauge.called)
+
+        async_set_status(
+            treatment_plan_pk=self.plan.pk,
+            status=TreatmentPlanStatus.SUCCESS,
+            start=False,
+        )
+        self.plan.refresh_from_db()
+        self.assertEquals(self.plan.status, TreatmentPlanStatus.SUCCESS)
+        self.assertIsNotNone(self.plan.started_at)
+        self.assertIsNotNone(self.plan.finished_at)
+        self.assertTrue(metrics_mock.gauge.called)