Fix flaky TeX Chromatic snapshots for Matcher and Sorter widgets

.changeset/stale-carrots-tease.md

@@ -0,0 +1,5 @@
+---
+"@khanacademy/perseus": patch
+---
+
+Fix flakey TeX Chromatic snapshots for Matcher and Sorter widgets.

packages/perseus/src/components/__tests__/with-asset-context.test.tsx

@@ -0,0 +1,94 @@
+import {render, screen} from "@testing-library/react";
+import * as React from "react";
+
+import AssetContext from "../../asset-context";
+import {withAssetContext} from "../with-asset-context";
+
+// Minimal test-only component used to exercise withAssetContext. It accepts
+// the props the HOC injects (`setAssetStatus` and `assetKey`) and surfaces
+// them in two ways the tests can observe: as a `data-asset-key` attribute on
+// the rendered button (for inspecting which key the HOC generated), and via
+// a click handler that calls `setAssetStatus` (so tests can verify the
+// injected function reaches the wrapped component and is callable).
+type MockComponentProps = {
+    label: string;
+    setAssetStatus: (assetKey: string, loaded: boolean) => void;
+    assetKey: string;
+};
+
+function MockComponent(props: MockComponentProps): React.ReactElement {
+    return (
+        <button
+            onClick={() => props.setAssetStatus(props.assetKey, true)}
+            data-testid={props.label}
+            data-asset-key={props.assetKey}
+        >
+            {props.label}
+        </button>
+    );
+}
+
+describe("withAssetContext", () => {
+    it("injects setAssetStatus from AssetContext as a prop", () => {
+        // Arrange
+        const setAssetStatusSpy = jest.fn();
+        const Wrapped = withAssetContext(MockComponent, "widget");
+
+        // Act — clicking calls setAssetStatus via the injected prop, which
+        // proves the function from the Provider reached MockComponent.
+        render(
+            <AssetContext.Provider
+                value={{
+                    assetStatuses: {},
+                    setAssetStatus: setAssetStatusSpy,
+                }}
+            >
+                <Wrapped label="hello" />
+            </AssetContext.Provider>,
+        );
+        screen.getByTestId("hello").click();
+
+        // Assert
+        expect(setAssetStatusSpy).toHaveBeenCalled();
+    });
+
+    it("injects an assetKey with the configured prefix", () => {
+        // Arrange, Act
+        const Wrapped = withAssetContext(MockComponent, "widget");
+        render(<Wrapped label="hello" />);
+
+        // Assert
+        expect(screen.getByTestId("hello").dataset.assetKey).toMatch(
+            /^widget-/,
+        );
+    });
+
+    it("generates a unique assetKey per rendered instance", () => {
+        // Arrange, Act
+        const Wrapped = withAssetContext(MockComponent, "widget");
+        render(
+            <>
+                <Wrapped label="first" />
+                <Wrapped label="second" />
+            </>,
+        );
+
+        // Assert
+        const first = screen.getByTestId("first");
+        const second = screen.getByTestId("second");
+        expect(first.dataset.assetKey).not.toBe(second.dataset.assetKey);
+    });
+
+    it("uses the AssetContext default no-op when no Provider is above", () => {
+        // Arrange, Act — render without an AssetContext.Provider.
+        const Wrapped = withAssetContext(MockComponent, "widget");
+        render(<Wrapped label="hello" />);
+
+        // Assert — the HOC still generates an assetKey (useId works without
+        // a Provider) and the default no-op setAssetStatus is callable (so
+        // the click handler doesn't crash).
+        const button = screen.getByTestId("hello");
+        expect(button.dataset.assetKey).toMatch(/^widget-/);
+        expect(() => button.click()).not.toThrow();
+    });
+});

packages/perseus/src/components/with-asset-context.tsx

@@ -0,0 +1,43 @@
+import * as React from "react";
+
+import AssetContext from "../asset-context";
+
+type SetAssetStatus = (assetKey: string, loaded: boolean) => void;
+
+type InjectedProps = {
+    setAssetStatus: SetAssetStatus;
+    assetKey: string;
+};
+
+export function withAssetContext<P extends InjectedProps>(
+    WrappedComponent: React.ComponentType<P>,
+    keyPrefix: string,
+): React.ForwardRefExoticComponent<
+    React.PropsWithoutRef<Omit<P, keyof InjectedProps>> &
+        React.RefAttributes<any>
+> {
+    const WithAssetContextComponent = React.forwardRef<
+        any,
+        Omit<P, keyof InjectedProps>
+    >((props, ref) => {
+        const {setAssetStatus} = React.useContext(AssetContext);
+        const uniqueId = React.useId();
+        const assetKey = `${keyPrefix}-${uniqueId}`;
+
+        return (
+            <WrappedComponent
+                // eslint-disable-next-line no-restricted-syntax -- TypeScript can't verify that spreading `Omit<P, K>` + adding `K`-properties reconstitutes P (the "P could be a different subtype" limitation). The runtime behavior is sound: we provide every field P needs.
+                {...(props as P)}
+                setAssetStatus={setAssetStatus}
+                assetKey={assetKey}
+                ref={ref}
+            />
+        );
+    });
+
+    WithAssetContextComponent.displayName = `withAssetContext(${
+        WrappedComponent.displayName || WrappedComponent.name || "Component"
+    })`;
+
+    return WithAssetContextComponent;
+}

packages/perseus/src/widgets/matcher/__docs__/matcher-initial-state-regression.stories.tsx

@@ -1,4 +1,12 @@
+import {
+    generateMatcherWidget,
+    generateTestPerseusItem,
+    generateTestPerseusRenderer,
+} from "@khanacademy/perseus-core";
+import * as React from "react";
+
 import {themeModes} from "../../../../../../.storybook/modes";
+import {ServerItemRendererWithDebugUI} from "../../../testing/server-item-renderer-with-debug-ui";
 
 import {matcherRendererDecorator} from "./matcher-renderer-decorator";
 
@@ -57,30 +65,39 @@ export const WithoutLabels: Story = {
     },
 };
 
-// TODO: Test flaky, commenting until it's fixed
-// Verifies TeX rendering in both column labels and item cards. The 500ms delay
-// gives MathJax time to finish rendering and onMeasure to settle before
-// Chromatic takes its snapshot.
-// export const WithTeX: Story = {
-//     decorators: [matcherRendererDecorator],
-//     parameters: {
-//         chromatic: {delay: 1000},
-//     },
-//     args: {
-//         ...sharedArgs,
-//         labels: ["$f(x)$", "$f'(x)$"],
-//         left: [
-//             "$f(x) = \\dfrac{1}{x}$",
-//             "$f(x) = \\dfrac{1}{x^2}$",
-//             "$f(x) = \\dfrac{1}{x^3}$",
-//         ],
-//         right: [
-//             "$f'(x) = -\\dfrac{1}{x^2}$",
-//             "$f'(x) = -\\dfrac{2}{x^3}$",
-//             "$f'(x) = -\\dfrac{3}{x^4}$",
-//         ],
-//     },
-// };
+// Verifies TeX rendering in both column labels and item cards. Rendered
+// through ServerItemRenderer so AssetContext.Provider is in the tree and
+// the measurement-cascade settling signal reaches the renderer.
+export const WithTeX: Story = {
+    render: () => (
+        <ServerItemRendererWithDebugUI
+            item={generateTestPerseusItem({
+                question: generateTestPerseusRenderer({
+                    content: "[[☃ matcher 1]]",
+                    widgets: {
+                        "matcher 1": generateMatcherWidget({
+                            options: {
+                                labels: ["$f(x)$", "$f'(x)$"],
+                                left: [
+                                    "$f(x) = \\dfrac{1}{x}$",
+                                    "$f(x) = \\dfrac{1}{x^2}$",
+                                    "$f(x) = \\dfrac{1}{x^3}$",
+                                ],
+                                right: [
+                                    "$f'(x) = -\\dfrac{1}{x^2}$",
+                                    "$f'(x) = -\\dfrac{2}{x^3}$",
+                                    "$f'(x) = -\\dfrac{3}{x^4}$",
+                                ],
+                                orderMatters: false,
+                                padding: true,
+                            },
+                        }),
+                    },
+                }),
+            })}
+        />
+    ),
+};
 
 // Verifies the orderMatters state: both columns render as draggable cards
 // (white background, visible border) rather than the left column appearing