From 6f4da6ea0f3dee1df6b1f5edd571ba778377fab6 Mon Sep 17 00:00:00 2001
From: Abigail Alexander <abigail.alexander@diamond.ac.uk>
Date: Thu, 25 Jun 2026 10:01:32 +0100
Subject: [PATCH] Add unique key to force img remount

---
 .../widgets/DynamicImage/demoImage.test.tsx   | 61 +++++++++++--------
 src/ui/widgets/DynamicImage/demoImage.tsx     |  7 ++-
 2 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/src/ui/widgets/DynamicImage/demoImage.test.tsx b/src/ui/widgets/DynamicImage/demoImage.test.tsx
index 5e718f2c..bd84e82f 100644
--- a/src/ui/widgets/DynamicImage/demoImage.test.tsx
+++ b/src/ui/widgets/DynamicImage/demoImage.test.tsx
@@ -1,5 +1,11 @@
 import React from "react";
-import { fireEvent, render, screen } from "@testing-library/react";
+import {
+  act,
+  fireEvent,
+  render,
+  screen,
+  waitFor
+} from "@testing-library/react";
 import { DemoImageComponent, buildMjpgPvUrls } from "./demoImage";
 import { newColor } from "../../../types/color";
 import { vi } from "vitest";
@@ -84,41 +90,39 @@ describe("DemoImageComponent", () => {
       />
     );
 
-    const img = screen.getByRole("img");
-
-    expect(img).toHaveAttribute("src", "http://a/TEST:PV");
+    expect(screen.getByRole("img")).toHaveAttribute("src", "http://a/TEST:PV");
 
-    fireEvent.error(img);
-    expect(img).toHaveAttribute("src", "http://b/TEST:PV");
+    fireEvent.error(screen.getByRole("img"));
+    expect(screen.getByRole("img")).toHaveAttribute("src", "http://b/TEST:PV");
 
-    fireEvent.error(img);
-    expect(img).toHaveAttribute("src", "http://c/TEST:PV");
+    fireEvent.error(screen.getByRole("img"));
+    expect(screen.getByRole("img")).toHaveAttribute("src", "http://c/TEST:PV");
   });
 
-  it("shows warning when all URLs fail", () => {
+  it("shows warning when all URLs fail", async () => {
     const data = [
       { value: newDType({ stringValue: "ignored" }, DAlarmNONE()) } as PvDatum
     ];
 
-    render(
-      <DemoImageComponent
-        macros={{}}
-        pvData={data}
-        mjpgEndpoints={["http://a", "http://b"]}
-      />
-    );
-
-    const img = screen.getByRole("img");
+    const { getByRole } = await act(() => {
+      return render(
+        <DemoImageComponent
+          macros={{}}
+          pvData={data}
+          mjpgEndpoints={["http://a", "http://b"]}
+        />
+      );
+    });
 
-    fireEvent.error(img); // move to second URL
-    fireEvent.error(img); // no more URLs → warning
+    fireEvent.error(getByRole("img")); // move to second URL
+    fireEvent.error(getByRole("img")); // no more URLs → warning
 
     expect(showWarningMock).toHaveBeenCalledWith(
       "Could not load mjpg image stream for the PV: TEST:PV"
     );
   });
 
-  it("resets failure counter after exhausting URLs", () => {
+  it("resets failure counter after exhausting URLs", async () => {
     const data = [
       { value: newDType({ stringValue: "ignored" }, DAlarmNONE()) } as PvDatum
     ];
@@ -131,15 +135,18 @@ describe("DemoImageComponent", () => {
       />
     );
 
-    const img = screen.getByRole("img");
-
-    fireEvent.error(img); // → b
-    fireEvent.error(img); // → warning + reset
+    fireEvent.error(screen.getByRole("img")); // → b
+    fireEvent.error(screen.getByRole("img")); // → warning + reset
 
     // Trigger again → should start sequence again from second URL
-    fireEvent.error(img);
+    fireEvent.error(screen.getByRole("img"));
 
-    expect(img).toHaveAttribute("src", "http://b/TEST:PV");
+    await waitFor(() => {
+      expect(screen.getByRole("img")).toHaveAttribute(
+        "src",
+        "http://b/TEST:PV"
+      );
+    });
   });
 
   it("handles missing endpoints without crashing", () => {
diff --git a/src/ui/widgets/DynamicImage/demoImage.tsx b/src/ui/widgets/DynamicImage/demoImage.tsx
index ce6e1600..7aba4f3a 100644
--- a/src/ui/widgets/DynamicImage/demoImage.tsx
+++ b/src/ui/widgets/DynamicImage/demoImage.tsx
@@ -1,4 +1,4 @@
-import React, { useState } from "react";
+import React, { useEffect, useState } from "react";
 import { Widget } from "../widget";
 import { PVComponent, PVWidgetPropType } from "../widgetProps";
 import {
@@ -86,9 +86,14 @@ export const DemoImageComponent = (
     }
   };
 
+  useEffect(() => {
+    setSrc(urls?.[numberOfFailures]);
+  }, [urls, numberOfFailures]);
+
   return (
     <Box
       component="img"
+      key={`image_${crypto.randomUUID()}`}
       src={src}
       alt={`PvName: ${effectivePvName}`}
       onError={handleError}