Merge remote-tracking branch 'origin/enterprise' into misc-fixes

testgen-ci-bot · testgen-ci-bot · commit 4f0134b5115e · 2026-06-08T14:27:56.000Z
diff --git a/testgen/commands/job_runner.py b/testgen/commands/job_runner.py
@@ -73,12 +73,12 @@ def _print_run_summary(job_id: UUID, job_key: str) -> None:
             case "run-profile":
                 run = session.scalars(select(ProfilingRun).where(ProfilingRun.job_execution_id == job_id)).first()
                 if run:
-                    status_msg = "Profiling encountered an error. Check log for details." if run.status == "Error" else "Profiling completed."
+                    status_msg = "Profiling encountered an error. Check log for details." if run.job_execution.status == JobStatus.ERROR else "Profiling completed."
                     click.echo(f"\n        {status_msg}\n        Run ID: {run.id}\n    ")
             case "run-tests" | "run-monitors":
                 run = session.scalars(select(TestRun).where(TestRun.job_execution_id == job_id)).first()
                 if run:
-                    status_msg = "Test execution encountered an error. Check log for details." if run.status == "Error" else "Test execution completed."
+                    status_msg = "Test execution encountered an error. Check log for details." if run.job_execution.status == JobStatus.ERROR else "Test execution completed."
                     click.echo(f"\n        {status_msg}\n        Run ID: {run.id}\n    ")
             case "run-test-generation":
                 click.echo("Test generation completed.")
diff --git a/testgen/common/models/profiling_run.py b/testgen/common/models/profiling_run.py
@@ -6,7 +6,7 @@
 
 from sqlalchemy import BigInteger, Column, Float, Integer, String, desc, func, select, text, update
 from sqlalchemy.dialects import postgresql
-from sqlalchemy.orm import InstrumentedAttribute
+from sqlalchemy.orm import InstrumentedAttribute, foreign, relationship
 from sqlalchemy.orm.attributes import flag_modified
 from sqlalchemy.sql.expression import case
 
@@ -117,6 +117,13 @@ class ProfilingRun(Entity):
     process_id: int = Column(Integer)
     job_execution_id: UUID | None = Column(postgresql.UUID(as_uuid=True), nullable=True)
 
+    job_execution = relationship(
+        JobExecution,
+        primaryjoin=foreign(job_execution_id) == JobExecution.id,
+        uselist=False,
+        viewonly=True,
+    )
+
     _default_order_by = (desc(profiling_starttime),)
     _minimal_columns = (
         id,
diff --git a/testgen/common/models/test_run.py b/testgen/common/models/test_run.py
@@ -5,6 +5,7 @@
 
 from sqlalchemy import BigInteger, Column, Float, ForeignKey, Integer, String, Text, desc, func, select, text, update
 from sqlalchemy.dialects import postgresql
+from sqlalchemy.orm import foreign, relationship
 from sqlalchemy.orm.attributes import flag_modified
 from sqlalchemy.sql.expression import case
 
@@ -121,6 +122,13 @@ class TestRun(Entity):
     process_id: int = Column(Integer)
     job_execution_id: UUID | None = Column(postgresql.UUID(as_uuid=True), nullable=True)
 
+    job_execution = relationship(
+        JobExecution,
+        primaryjoin=foreign(job_execution_id) == JobExecution.id,
+        uselist=False,
+        viewonly=True,
+    )
+
     _default_order_by = (desc(test_starttime),)
     _minimal_columns = (
         id,
diff --git a/testgen/common/notifications/monitor_run.py b/testgen/common/notifications/monitor_run.py
@@ -236,7 +236,7 @@ def send_monitor_notifications(test_run: TestRun, result_list_ct=20):
                 notification.recipients,
                 {
                     "summary": {
-                        "test_endtime": test_run.test_endtime,
+                        "test_endtime": test_run.job_execution.completed_at,
                         "table_groups_name": table_group.table_groups_name,
                         "project_name": project.project_name,
                         "table_name": table_name,
diff --git a/testgen/common/notifications/notifications.py b/testgen/common/notifications/notifications.py
@@ -1,8 +1,6 @@
 import math
 from datetime import datetime
 
-from testgen.common.models.profiling_run import ProfilingRunStatus
-from testgen.common.models.test_definition import TestRunStatus
 from testgen.common.notifications.base import BaseEmailTemplate
 from testgen.utils import friendly_score
 
@@ -29,12 +27,6 @@ def format_duration_helper(self, start_time: datetime, end_time: datetime) -> st
         formatted = " ".join([ f"{unit[0]}{unit[1]}" for unit in units if unit[0] ])
         return formatted.strip() or "< 1s"
 
-    def format_status_helper(self, status: TestRunStatus | ProfilingRunStatus) -> str:
-        return {
-            "Complete": "Completed",
-            "Cancelled": "Canceled",
-        }.get(status, status)
-
     def format_score_helper(self, score: float) -> str:
         return friendly_score(score)
 
diff --git a/testgen/common/notifications/profiling_run.py b/testgen/common/notifications/profiling_run.py
@@ -4,6 +4,7 @@
 from sqlalchemy import select
 
 from testgen import settings
+from testgen.common.enums import JOB_STATUS_LABEL, JobStatus
 from testgen.common.models import get_current_session, with_database_session
 from testgen.common.models.hygiene_issue import HygieneIssue
 from testgen.common.models.notification_settings import (
@@ -23,7 +24,7 @@ class ProfilingRunEmailTemplate(BaseNotificationTemplate):
 
     def get_subject_template(self) -> str:
         return (
-            "[TestGen] Profiling Run {{format_status profiling_run.status}}: {{table_groups_name}}"
+            "[TestGen] Profiling Run {{profiling_run.status_label}}: {{table_groups_name}}"
             "{{#if issue_count}}"
             ' | {{format_number issue_count}} hygiene {{pluralize issue_count "issue" "issues"}}'
             "{{/if}}"
@@ -32,9 +33,9 @@ def get_subject_template(self) -> str:
     def get_title_template(self):
         return """
           TestGen Profiling Run - <span class="
-          {{#if (eq profiling_run.status 'Error')}} text-red {{/if}}
-          {{#if (eq profiling_run.status 'Cancelled')}} text-purple {{/if}}
-          ">{{format_status profiling_run.status}}</span>
+          {{#if (eq profiling_run.status 'error')}} text-red {{/if}}
+          {{#if (eq profiling_run.status 'canceled')}} text-purple {{/if}}
+          ">{{profiling_run.status_label}}</span>
         """
 
     def get_main_content_template(self):
@@ -85,7 +86,7 @@ def get_main_content_template(self):
                 </tr>
                 <tr>
                   <td class="summary__subtitle">
-                    {{#if (eq profiling_run.status 'Complete')}}
+                    {{#if (eq profiling_run.status 'completed')}}
                         {{#if (eq notification_trigger 'on_changes')}}
                             Profiling run detected new hygiene issues.
                         {{/if}}
@@ -95,15 +96,15 @@ def get_main_content_template(self):
                             {{/if}}
                         {{/if}}
                     {{/if}}
-                    {{#if (eq profiling_run.status 'Error')}}
+                    {{#if (eq profiling_run.status 'error')}}
                         Profiling encountered an error.
                     {{/if}}
-                    {{#if (eq profiling_run.status 'Cancelled')}}
+                    {{#if (eq profiling_run.status 'canceled')}}
                         Profiling run was canceled.
                     {{/if}}
                   </td>
                 </tr>
-                {{#if (eq profiling_run.status 'Complete')}}
+                {{#if (eq profiling_run.status 'completed')}}
                 <tr>
                   <td colspan="2" style="padding-top: 12px; padding-bottom: 12px;">
                     <table cellspacing="0" cellpadding="0">
@@ -131,12 +132,12 @@ def get_main_content_template(self):
                   </td>
                 </tr>
                 {{/if}}
-                {{#if (eq profiling_run.status 'Error')}}
+                {{#if (eq profiling_run.status 'error')}}
                 <tr>
                   <td><div class="code">{{profiling_run.log_message}}</div></td>
                 </tr>
                 {{/if}}
-                {{#if (eq profiling_run.status 'Complete')}}
+                {{#if (eq profiling_run.status 'completed')}}
                 <tr>
                   <td>
                     <a class="link" href="{{profiling_run.issues_url}}" target="_blank">View {{format_number issue_count}} issues &gt;</a>
@@ -241,6 +242,8 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct
     if not notifications:
         return
 
+    job_execution = profiling_run.job_execution
+
     previous_run = profiling_run.get_previous()
     issues = list(
         HygieneIssue.select_with_diff(
@@ -251,7 +254,7 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct
     )
 
     triggers = {ProfilingRunNotificationTrigger.always}
-    if profiling_run.status in ("Error", "Cancelled") or {None, True} & {is_new for _, is_new in issues}:
+    if job_execution.status in (JobStatus.ERROR, JobStatus.CANCELED) or {None, True} & {is_new for _, is_new in issues}:
         triggers.add(ProfilingRunNotificationTrigger.on_changes)
 
     notifications = [ns for ns in notifications if ns.trigger in triggers]
@@ -320,10 +323,11 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct
                     "&source=email"
                 )
             ),
-            "start_time": profiling_run.profiling_starttime,
-            "end_time": profiling_run.profiling_endtime,
-            "status": profiling_run.status,
-            "log_message": profiling_run.log_message,
+            "start_time": job_execution.started_at,
+            "end_time": job_execution.completed_at,
+            "status": job_execution.status,
+            "status_label": JOB_STATUS_LABEL.get(job_execution.status, job_execution.status),
+            "log_message": job_execution.error_message,
             "table_ct": profiling_run.table_ct,
             "column_ct": profiling_run.column_ct,
         },
diff --git a/testgen/common/notifications/test_run.py b/testgen/common/notifications/test_run.py
@@ -3,6 +3,7 @@
 from sqlalchemy import case, literal, select
 
 from testgen import settings
+from testgen.common.enums import JobStatus
 from testgen.common.models import get_current_session, with_database_session
 from testgen.common.models.notification_settings import (
     TestRunNotificationSettings,
@@ -261,7 +262,7 @@ def send_test_run_notifications(test_run: TestRun, result_list_ct=20, result_sta
                 changed_td_id_list.extend(td_id_list)
 
     triggers = {TestRunNotificationTrigger.always}
-    if test_run.status in ("Error", "Cancelled"):
+    if test_run.job_execution.status in (JobStatus.ERROR, JobStatus.CANCELED):
         triggers.update(TestRunNotificationTrigger)
     else:
         if test_run.error_ct + test_run.failed_ct:
diff --git a/tests/unit/common/notifications/test_monitor_run_notifications.py b/tests/unit/common/notifications/test_monitor_run_notifications.py
@@ -29,6 +29,12 @@ def create_test_result(table_name, test_type, message, result_code=0):
     return mock
 
 
+def make_monitor_run(completed_at="2024-01-15T10:30:00Z"):
+    run = TestRun(id="monitor-run-id", test_suite_id="monitor-suite-id")
+    run.job_execution = Mock(completed_at=completed_at)
+    return run
+
+
 @pytest.fixture
 def ns_select_result():
     return [
@@ -110,11 +116,7 @@ def test_send_monitor_notifications(
         send_mock,
         persisted_setting_mock,
 ):
-    test_run = TestRun(
-        id="monitor-run-id",
-        test_suite_id="monitor-suite-id",
-        test_endtime="2024-01-15T10:30:00Z",
-    )
+    test_run = make_monitor_run()
 
     table_group = Mock(spec=TableGroup)
     table_group.id = "tg-id"
@@ -193,11 +195,7 @@ def test_send_monitor_notifications_early_exit(
         test_result_select_where_mock,
         send_mock,
 ):
-    test_run = TestRun(
-        id="monitor-run-id",
-        test_suite_id="monitor-suite-id",
-        test_endtime="2024-01-15T10:30:00Z",
-    )
+    test_run = make_monitor_run()
 
     if not has_notifications:
         ns_select_patched.return_value = []
@@ -219,11 +217,7 @@ def test_send_monitor_notifications_anomaly_counts(
         send_mock,
         persisted_setting_mock,
 ):
-    test_run = TestRun(
-        id="monitor-run-id",
-        test_suite_id="monitor-suite-id",
-        test_endtime="2024-01-15T10:30:00Z",
-    )
+    test_run = make_monitor_run()
 
     table_group = Mock(spec=TableGroup)
     table_group.id = "tg-id"
@@ -270,11 +264,7 @@ def test_send_monitor_notifications_url_construction(
         send_mock,
         persisted_setting_mock,
 ):
-    test_run = TestRun(
-        id="monitor-run-id",
-        test_suite_id="monitor-suite-id",
-        test_endtime="2024-01-15T10:30:00Z",
-    )
+    test_run = make_monitor_run()
 
     table_group = Mock(spec=TableGroup)
     table_group.id = "tg-123"
diff --git a/tests/unit/common/notifications/test_profiling_run_notifications.py b/tests/unit/common/notifications/test_profiling_run_notifications.py
@@ -4,6 +4,7 @@
 
 import pytest
 
+from testgen.common.enums import JOB_STATUS_LABEL, JobStatus
 from testgen.common.models.hygiene_issue import IssueCount
 from testgen.common.models.notification_settings import (
     ProfilingRunNotificationSettings,
@@ -15,6 +16,10 @@
 pytestmark = pytest.mark.unit
 
 
+def make_job_execution(status, started_at=None, completed_at=None, error_message=None):
+    return Mock(status=status, started_at=started_at, completed_at=completed_at, error_message=error_message)
+
+
 def create_ns(**kwargs):
     with patch("testgen.common.notifications.profiling_run.ProfilingRunNotificationSettings.save"):
         return ProfilingRunNotificationSettings.create("proj", None, **kwargs)
@@ -79,12 +84,12 @@ def get_prev_mock():
 @pytest.mark.parametrize(
     ("profiling_run_status", "has_prev_run", "issue_count", "new_issue_count", "expected_triggers"),
     (
-        ("Error", True, 25, 0, ("always", "on_changes")),
-        ("Error", True, 0, 0, ("always", "on_changes")),
-        ("Cancelled", True, 50, 10, ("always", "on_changes")),
-        ("Complete", True, 50, 10, ("always", "on_changes")),
-        ("Complete", True, 15, 0, ("always",)),
-        ("Complete", False, 15, 15, ("always", "on_changes")),
+        (JobStatus.ERROR, True, 25, 0, ("always", "on_changes")),
+        (JobStatus.ERROR, True, 0, 0, ("always", "on_changes")),
+        (JobStatus.CANCELED, True, 50, 10, ("always", "on_changes")),
+        (JobStatus.COMPLETED, True, 50, 10, ("always", "on_changes")),
+        (JobStatus.COMPLETED, True, 15, 0, ("always",)),
+        (JobStatus.COMPLETED, False, 15, 15, ("always", "on_changes")),
     ),
 )
 def test_send_profiling_run_notification(
@@ -104,9 +109,9 @@ def test_send_profiling_run_notification(
         id="pr-id",
         job_execution_id="pr-id",
         table_groups_id="tg-id",
-        status=profiling_run_status,
         project_code="proj",
     )
+    profiling_run.job_execution = make_job_execution(profiling_run_status)
     get_prev_mock.return_value = ProfilingRun(id="pr-prev-id") if has_prev_run else None
     new_count = iter(count())
     priorities = ("Definite", "Likely", "Possible", "High", "Moderate")
@@ -144,6 +149,7 @@ def test_send_profiling_run_notification(
                         "start_time": None,
                         "end_time": None,
                         "status": profiling_run_status,
+                        "status_label": JOB_STATUS_LABEL[profiling_run_status],
                         "log_message": None,
                         "table_ct": None,
                         "column_ct": None,
diff --git a/tests/unit/common/notifications/test_test_run_notifications.py b/tests/unit/common/notifications/test_test_run_notifications.py
diff --git a/tests/unit/common/test_enums.py b/tests/unit/common/test_enums.py

Original file line number	Diff line number	Diff line change
`@@ -236,7 +236,7 @@ def send_monitor_notifications(test_run: TestRun, result_list_ct=20):`
`236`	`236`	`notification.recipients,`
`237`	`237`	`{`
`238`	`238`	`"summary": {`
`239`		`- "test_endtime": test_run.test_endtime,`
	`239`	`+ "test_endtime": test_run.job_execution.completed_at,`
`240`	`240`	`"table_groups_name": table_group.table_groups_name,`
`241`	`241`	`"project_name": project.project_name,`
`242`	`242`	`"table_name": table_name,`