From 502a4ad4a0df411e3c8bcaec7e08c8582bf3bc6f Mon Sep 17 00:00:00 2001
From: Ryan Britton <ryan.britton@canonical.com>
Date: Mon, 13 Jan 2025 20:42:41 +0000
Subject: [PATCH] Allow end test by test execution id

---
 .../controllers/test_executions/end_test.py   |  9 +++-
 .../controllers/test_executions/models.py     | 10 +++-
 .../test_executions/test_end_test.py          | 47 +++++++++++++++++++
 3 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/backend/test_observer/controllers/test_executions/end_test.py b/backend/test_observer/controllers/test_executions/end_test.py
index 82a780ba..61458266 100644
--- a/backend/test_observer/controllers/test_executions/end_test.py
+++ b/backend/test_observer/controllers/test_executions/end_test.py
@@ -15,7 +15,7 @@
 #
 
 from fastapi import APIRouter, Depends, HTTPException
-from sqlalchemy import select
+from sqlalchemy import select, or_
 from sqlalchemy.orm import Session, joinedload
 
 from test_observer.data_access.models import (
@@ -67,7 +67,12 @@ def _find_related_test_execution(
     return (
         db.execute(
             select(TestExecution)
-            .where(TestExecution.ci_link == request.ci_link)
+            .where(
+                or_(
+                    TestExecution.ci_link == request.ci_link,
+                    TestExecution.id == request.test_execution_id,
+                )
+            )
             .options(
                 joinedload(TestExecution.artefact_build).joinedload(
                     ArtefactBuild.artefact
diff --git a/backend/test_observer/controllers/test_executions/models.py b/backend/test_observer/controllers/test_executions/models.py
index 64416eeb..03d3f720 100644
--- a/backend/test_observer/controllers/test_executions/models.py
+++ b/backend/test_observer/controllers/test_executions/models.py
@@ -29,6 +29,7 @@
     Field,
     HttpUrl,
     field_validator,
+    model_validator,
 )
 
 from test_observer.common.constants import PREVIOUS_TEST_RESULT_COUNT
@@ -102,11 +103,18 @@ class TestResultRequest(BaseModel):
 
 
 class EndTestExecutionRequest(BaseModel):
-    ci_link: Annotated[str, HttpUrl]
+    ci_link: Annotated[str, HttpUrl] | None = None
+    test_execution_id: int | None = None
     c3_link: Annotated[str, HttpUrl] | None = None
     checkbox_version: str | None = None
     test_results: list[C3TestResult]
 
+    @model_validator(mode="after")
+    def ensure_test_execution_identifier(self) -> "EndTestExecutionRequest":
+        if sum([self.ci_link is None, self.test_execution_id is None]) != 1:
+            raise ValueError("Provide exactly one of 'ci_link' or 'test_execution_id'")
+        return self
+
 
 class TestExecutionsPatchRequest(BaseModel):
     __test__ = False
diff --git a/backend/tests/controllers/test_executions/test_end_test.py b/backend/tests/controllers/test_executions/test_end_test.py
index f3b3fde4..94fbecdb 100644
--- a/backend/tests/controllers/test_executions/test_end_test.py
+++ b/backend/tests/controllers/test_executions/test_end_test.py
@@ -138,3 +138,50 @@ def test_end_test_updates_template_id(
 
     assert response.status_code == 200
     assert test_execution.test_results[0].test_case.template_id == "some template id"
+
+
+def test_end_test_no_test_identifier(test_client: TestClient):
+    response = test_client.put(
+        "/v1/test-executions/end-test",
+        json={
+            "test_results": [],
+        },
+    )
+
+    assert response.status_code == 422
+
+
+def test_end_test_by_test_execution_id(
+    test_client: TestClient, generator: DataGenerator
+):
+    artefact = generator.gen_artefact(StageName.beta)
+    artefact_build = generator.gen_artefact_build(artefact)
+    environment = generator.gen_environment()
+    test_execution = generator.gen_test_execution(
+        artefact_build, environment, ci_link="http://localhost"
+    )
+    generator.gen_artefact_build_environment_review(artefact_build, environment)
+
+    response = test_client.put(
+        "/v1/test-executions/end-test",
+        json={
+            "test_execution_id": test_execution.id,
+            "test_results": [],
+        },
+    )
+
+    assert response.status_code == 200
+    assert test_execution.status == TestExecutionStatus.PASSED
+
+
+def test_end_test_multiple_identifiers(test_client: TestClient):
+    response = test_client.put(
+        "/v1/test-executions/end-test",
+        json={
+            "ci_link": "http://localhost",
+            "test_execution_id": 1,
+            "test_results": [],
+        },
+    )
+
+    assert response.status_code == 422