feat(nimbus): create update leading branch image forms

Author: moibra05

experimenter/experimenter/nimbus_ui/forms.py

@@ -50,6 +50,10 @@ def get_changelog_message(self) -> str:
 
     def save(self, *args, **kwargs):
         experiment = super().save(*args, **kwargs)
+
+        if type(experiment) is NimbusBranchScreenshot:
+            experiment = self.instance.branch.experiment
+
         generate_nimbus_changelog(
             experiment, self.request.user, self.get_changelog_message()
         )
@@ -1771,3 +1775,24 @@ def save(self, commit=True):
 
     def get_changelog_message(self):
         return f"{self.request.user} updated collaborators"
+
+
+class BranchLeadingScreenshotForm(NimbusChangeLogFormMixin, forms.ModelForm):
+    image = forms.ImageField(
+        required=False,
+        widget=forms.FileInput(attrs={"class": "form-control"}),
+    )
+    description = forms.CharField(
+        required=False,
+        widget=forms.TextInput(attrs={"class": "form-control"}),
+    )
+
+    class Meta:
+        model = NimbusBranchScreenshot
+        fields = ["image", "description"]
+
+    def get_changelog_message(self):
+        return (
+            f"{self.request.user} updated leading screenshot for "
+            f"{self.instance.branch.slug} branch"
+        )

experimenter/experimenter/nimbus_ui/templates/common/branch_card.html

@@ -1,17 +1,31 @@
 {% load nimbus_extras %}
 
-<div class="d-flex gap-4">
+<div class="d-flex gap-3">
   {% with shot=branch.screenshots.first %}
     {% if shot and shot.image %}
-      <img src="{{ shot.image.url }}"
-           alt="{{ shot.description|default:branch.name }}"
-           class="border border-secondary-subtle rounded w-50"
-           style="height: 120px;
-                  object-fit: cover" />
+      <div class="position-relative border border-secondary-subtle rounded w-50 h-100">
+        <img src="{{ shot.image.url }}"
+             alt="{{ shot.description|default:branch.name }}"
+             class="rounded w-100"
+             style="height: 120px;
+                    object-fit: cover" />
+        <button type="button"
+                data-bs-toggle="modal"
+                data-bs-target="#{{ experiment_slug }}-{{ branch.slug }}-image-edit"
+                class="position-absolute border-0 opacity-75 bg-secondary-subtle top-0 end-0 p-2 m-1 rounded lh-1">
+          <i class="fa-solid fa-up-right-and-down-left-from-center text-muted"></i>
+        </button>
+      </div>
     {% else %}
-      <div class="border border-secondary-subtle rounded text-center d-flex align-items-center justify-content-center w-50"
+      <div class="position-relative border border-secondary-subtle rounded text-center d-flex align-items-center justify-content-center w-50"
            style="height: 120px">
         <small class="text-muted">No screenshot</small>
+        <button type="button"
+                data-bs-toggle="modal"
+                data-bs-target="#{{ experiment_slug }}-{{ branch.slug }}-image-edit"
+                class="position-absolute border-0 opacity-75 bg-secondary-subtle top-0 end-0 p-2 m-1 rounded lh-1">
+          <i class="fa-solid fa-up-right-and-down-left-from-center text-muted"></i>
+        </button>
       </div>
     {% endif %}
   {% endwith %}
@@ -28,3 +42,63 @@ <h6 class="mt-2">{{ branch.name }}</h6>
     <p class="text-muted">{{ branch.description|truncatechars:100 }}</p>
   </div>
 </div>
+<div class="modal fade"
+     id="{{ experiment_slug }}-{{ branch.slug }}-image-edit"
+     tabindex="-1"
+     aria-labelledby="{{ experiment_slug }}-{{ branch.slug }}-image-edit"
+     aria-hidden="true">
+  <div class="modal-dialog modal-dialog-centered modal-lg">
+    <div class="modal-content p-4">
+      <div class="modal-header border-0">
+        <button type="button"
+                class="btn-close"
+                data-bs-dismiss="modal"
+                aria-label="Close"></button>
+      </div>
+      <div class="modal-body">
+        <form method="post"
+              id="branch-screenshot-form-{{ experiment_slug }}-{{ branch.slug }}"
+              hx-encoding="multipart/form-data"
+              hx-post="{% url "nimbus-ui-branch-leading-screenshot-upload" slug=experiment_slug branch_slug=branch.slug %}"
+              class="accordion">
+          {% csrf_token %}
+          {% if screenshot_form.instance.pk and screenshot_form.instance.image %}
+            <div class="position-relative border border-secondary-subtle rounded h-100 mb-3">
+              <img src="{{ screenshot_form.instance.image.url }}"
+                   alt="{{ screenshot_form.instance.description }}"
+                   class="w-100 rounded">
+            </div>
+          {% else %}
+            <div class="border border-secondary-subtle rounded mb-3">
+              <small class="text-muted d-block p-5 text-center">No screenshot uploaded</small>
+            </div>
+          {% endif %}
+          {{ screenshot_form.image }}
+          {% if screenshot_form.instance.image %}
+            <div class="accordion-item border border-1 rounded mt-4">
+              <button class="accordion-button shadow-none bg-transparent text-body"
+                      type="button"
+                      data-bs-toggle="collapse"
+                      data-bs-target="#screenshot-form-{{ experiment_slug }}-{{ branch.slug }}"
+                      aria-expanded="true"
+                      aria-controls="screenshot-form-{{ experiment_slug }}-{{ branch.slug }}">
+                <div class="d-flex flex-column align-items-start">
+                  <h6 class="mb-1">Add image description</h6>
+                  <small class="text-muted">Add a short description so people who can't see the image know what it shows.</small>
+                </div>
+              </button>
+              <div id="screenshot-form-{{ experiment_slug }}-{{ branch.slug }}"
+                   class="accordion-collapse collapse show">
+                <div class="accordion-body pt-0">{{ screenshot_form.description }}</div>
+              </div>
+            </div>
+          {% endif %}
+          <div class="d-flex justify-content-end gap-2 mt-4">
+            <button type="button" class="btn btn-secondary" data-bs-dismiss="modal">Close</button>
+            <button type="submit" class="btn btn-primary">Save changes</button>
+          </div>
+        </form>
+      </div>
+    </div>
+  </div>
+</div>

experimenter/experimenter/nimbus_ui/templates/nimbus_experiments/results-new-inner.html

@@ -89,8 +89,12 @@ <h5>Hypothesis</h5>
           <div class="row row-cols-2 row-cols-xxl-3 g-4">
             {% for branch in branch_data %}
               <div class="col">
-                {% include "common/branch_card.html" with branch=branch %}
+                {% for branch_slug, branch_leading_screenshot_form in branch_leading_screenshot_forms.items %}
+                  {% if branch_slug == branch.slug %}
+                    {% include "common/branch_card.html" with branch=branch experiment_slug=experiment.slug screenshot_form=branch_leading_screenshot_form %}
 
+                  {% endif %}
+                {% endfor %}
               </div>
             {% endfor %}
           </div>

experimenter/experimenter/nimbus_ui/tests/test_forms.py

@@ -41,6 +41,7 @@
     ApproveEndExperimentForm,
     ApproveUpdateRolloutForm,
     AudienceForm,
+    BranchLeadingScreenshotForm,
     BranchScreenshotCreateForm,
     BranchScreenshotDeleteForm,
     CancelEndEnrollmentForm,
@@ -4411,3 +4412,49 @@ def test_feature_unsubscribe_form_removes_subscriber(self):
         self.assertTrue(form.is_valid())
         form.save()
         self.assertNotIn(self.request.user, feature_config.subscribers.all())
+
+
+class ResultsEditBranchLeadingImageTests(RequestFormTestCase):
+    def test_edit_leading_image(self):
+        application = NimbusExperiment.Application.DESKTOP
+        experiment = NimbusExperimentFactory.create_with_lifecycle(
+            NimbusExperimentFactory.Lifecycles.CREATED,
+            application=application,
+        )
+        experiment.branches.all().delete()
+        experiment.changes.all().delete()
+
+        reference_branch = NimbusBranchFactory.create(experiment=experiment, ratio=1)
+        experiment.reference_branch = reference_branch
+        experiment.save()
+
+        reference_screenshot = reference_branch.screenshots.first()
+
+        image_bytes = io.BytesIO()
+        image = Image.new("RGB", (10, 10), color="red")
+        image.save(image_bytes, format="PNG")
+        image_bytes.seek(0)
+        dummy_image = SimpleUploadedFile(
+            "test.png", image_bytes.read(), content_type="image/png"
+        )
+
+        form = BranchLeadingScreenshotForm(
+            instance=reference_screenshot,
+            data={
+                "description": "Updated control screenshot",
+            },
+            files={
+                "image": dummy_image,
+            },
+            request=self.request,
+        )
+
+        form.save()
+        experiment.refresh_from_db()
+
+        self.assertEqual(
+            experiment.reference_branch.screenshots.get(
+                id=reference_screenshot.id
+            ).description,
+            "Updated control screenshot",
+        )

experimenter/experimenter/nimbus_ui/tests/test_views.py

@@ -1,11 +1,14 @@
 import datetime
+import io
 import json
 from unittest.mock import patch
 
 from django.conf import settings
+from django.core.files.uploadedfile import SimpleUploadedFile
 from django.test import TestCase
 from django.urls import reverse
 from parameterized import parameterized
+from PIL import Image
 
 from experimenter.base.tests.factories import (
     CountryFactory,
@@ -3028,6 +3031,62 @@ def test_get_redirects_to_next_step(self, current_url, next_url, data):
         )
 
 
+@mock_valid_outcomes
+class TestResultsEditBranchImagesView(AuthTestCase):
+    def test_upload_updates_screenshot(self):
+        experiment = NimbusExperimentFactory.create(owner=self.user)
+        branch = NimbusBranchFactory.create(experiment=experiment, slug="t1")
+
+        url = reverse(
+            "nimbus-ui-branch-leading-screenshot-upload",
+            kwargs={"slug": experiment.slug, "branch_slug": branch.slug},
+        )
+
+        image_bytes = io.BytesIO()
+        image = Image.new("RGB", (10, 10), color="red")
+        image.save(image_bytes, format="PNG")
+        image_bytes.seek(0)
+        dummy_image = SimpleUploadedFile(
+            "test.png", image_bytes.read(), content_type="image/png"
+        )
+
+        response = self.client.post(url, {"image": dummy_image})
+
+        self.assertEqual(response.status_code, 200)
+
+        branch.refresh_from_db()
+        shot = branch.screenshots.first()
+        self.assertIsNotNone(shot)
+
+    def test_upload_creates_screenshot(self):
+        experiment = NimbusExperimentFactory.create(owner=self.user)
+        experiment.branches.all().delete()
+
+        branch = NimbusBranchFactory.create(experiment=experiment, slug="t1")
+        branch.screenshots.all().delete()
+
+        url = reverse(
+            "nimbus-ui-branch-leading-screenshot-upload",
+            kwargs={"slug": experiment.slug, "branch_slug": branch.slug},
+        )
+
+        image_bytes = io.BytesIO()
+        image = Image.new("RGB", (10, 10), color="red")
+        image.save(image_bytes, format="PNG")
+        image_bytes.seek(0)
+        dummy_image = SimpleUploadedFile(
+            "test.png", image_bytes.read(), content_type="image/png"
+        )
+
+        response = self.client.post(url, {"image": dummy_image})
+
+        self.assertEqual(response.status_code, 200)
+
+        branch.refresh_from_db()
+        shot = branch.screenshots.first()
+        self.assertIsNotNone(shot)
+
+
 @mock_valid_outcomes
 class TestResultsView(AuthTestCase):
     def test_render_to_response(self):

experimenter/experimenter/nimbus_ui/urls.py

@@ -10,6 +10,7 @@
     BranchDeleteView,
     BranchesPartialUpdateView,
     BranchesUpdateView,
+    BranchLeadingScreenshotView,
     BranchScreenshotCreateView,
     BranchScreenshotDeleteView,
     CancelEndEnrollmentView,
@@ -298,4 +299,9 @@
         BranchScreenshotDeleteView.as_view(),
         name="nimbus-ui-delete-branch-screenshot",
     ),
+    re_path(
+        r"^(?P<slug>[\w-]+)/branches/(?P<branch_slug>[\w-]+)/leading-screenshot/",
+        BranchLeadingScreenshotView.as_view(),
+        name="nimbus-ui-branch-leading-screenshot-upload",
+    ),
 ]

experimenter/experimenter/nimbus_ui/views.py

@@ -36,6 +36,7 @@
     ApproveEndExperimentForm,
     ApproveUpdateRolloutForm,
     AudienceForm,
+    BranchLeadingScreenshotForm,
     BranchScreenshotCreateForm,
     BranchScreenshotDeleteForm,
     CancelEndEnrollmentForm,
@@ -671,6 +672,38 @@ class ApproveUpdateRolloutView(StatusUpdateView):
     form_class = ApproveUpdateRolloutForm
 
 
+class BranchLeadingScreenshotView(
+    NimbusExperimentViewMixin,
+    RequestFormMixin,
+    UpdateView,
+):
+    form_class = BranchLeadingScreenshotForm
+
+    def get_branch(self):
+        return (
+            self.get_object().branches.filter(slug=self.kwargs.get("branch_slug")).first()
+        )
+
+    def get_form_kwargs(self):
+        kwargs = super().get_form_kwargs()
+
+        if branch := self.get_branch():
+            if screenshot := branch.screenshots.first():
+                kwargs["instance"] = screenshot
+            else:
+                # create an unsaved placeholder instance of the correct model
+                screenshot_model = branch.screenshots.model
+                kwargs["instance"] = screenshot_model(branch=branch)
+
+        return kwargs
+
+    def form_valid(self, form):
+        form.save()
+        response = HttpResponse()
+        response.headers["HX-Refresh"] = "true"
+        return response
+
+
 class NewResultsView(NimbusExperimentViewMixin, DetailView):
     template_name = "nimbus_experiments/results-new.html"
 
@@ -710,6 +743,13 @@ def get_context_data(self, **kwargs):
             analysis_basis, selected_segment, selected_reference_branch
         )
 
+        branch_leading_screenshot_forms = {
+            branch.slug: BranchLeadingScreenshotForm(instance=branch.screenshots.first())
+            for branch in experiment.branches.all()
+        }
+
+        context["branch_leading_screenshot_forms"] = branch_leading_screenshot_forms
+
         relative_metric_changes = {}
 
         all_metrics = experiment.get_metric_data(