fix: copied code no longer includes markdown comments

src/components/MDXComponents/MDXCode.tsx

@@ -1,12 +1,11 @@
 import { useId, useState, useEffect } from 'react';
-import { CopyToClipboard } from 'react-copy-to-clipboard';
 import { Prism, Highlight } from 'prism-react-renderer';
 import { theme } from './code-theme';
-import { Button, Flex, View, VisuallyHidden } from '@aws-amplify/ui-react';
+import { Flex, View } from '@aws-amplify/ui-react';
 import { versions } from '@/constants/versions';
-import { trackCopyClicks } from '@/utils/track';
 import type { MDXCodeProps } from './types';
 import { TokenList } from './TokenList';
+import { MDXCopyCodeButton } from './MDXCopyCodeButton';
 
 (typeof global !== 'undefined' ? global : window).Prism = Prism;
 require('prismjs/components/prism-java');
@@ -39,21 +38,12 @@ export const MDXCode = ({
   testId,
   title
 }: MDXCodeProps) => {
-  const [copied, setCopied] = useState(false);
   const [code, setCode] = useState(codeString);
   const shouldShowCopy = language !== 'console';
   const shouldShowHeader = shouldShowCopy || title;
   const titleId = `${useId()}-titleID`;
   const codeId = `${useId()}-codeID`;
 
-  const copy = () => {
-    trackCopyClicks(codeString);
-    setCopied(true);
-    setTimeout(() => {
-      setCopied(false);
-    }, 2000);
-  };
-
   useEffect(() => {
     setCode(addVersions(codeString));
   }, [codeString]);
@@ -71,18 +61,11 @@ export const MDXCode = ({
                   </View>
                 ) : null}
                 {shouldShowCopy ? (
-                  <CopyToClipboard text={codeString} onCopy={copy}>
-                    <Button
-                      size="small"
-                      variation="link"
-                      disabled={copied}
-                      className="code-copy"
-                      aria-describedby={title ? undefined : codeId}
-                    >
-                      {copied ? 'Copied!' : 'Copy'}
-                      <VisuallyHidden>{title} code example</VisuallyHidden>
-                    </Button>
-                  </CopyToClipboard>
+                  <MDXCopyCodeButton
+                    codeString={codeString}
+                    title={title}
+                    codeId={codeId}
+                  />
                 ) : null}
               </Flex>
             ) : null}

src/components/MDXComponents/MDXCopyCodeButton.tsx

@@ -0,0 +1,61 @@
+import { useState } from 'react';
+import { CopyToClipboard } from 'react-copy-to-clipboard';
+import { Button, VisuallyHidden } from '@aws-amplify/ui-react';
+import { trackCopyClicks } from '@/utils/track';
+
+export const prepareCopyText = (codeString: string): string => {
+  // We need to strip out markdown comments from the code string
+  // so they don't show up in our copied text
+  const highlightStartText = /\/\/\s?highlight-start/g;
+  const highlightEndText = /\/\/\s?highlight-end/g;
+  const highlightNextLine = /\/\/\s?highlight-next-line/g;
+
+  return codeString
+    .replace(highlightStartText, '')
+    .replace(highlightEndText, '')
+    .replace(highlightNextLine, '');
+};
+
+interface MDXCopyCodeButtonProps {
+  codeId: string;
+  codeString: string;
+  testId?: string;
+  title?: string;
+}
+
+export const MDXCopyCodeButton = ({
+  codeString,
+  title,
+  codeId,
+  testId
+}: MDXCopyCodeButtonProps) => {
+  const [copied, setCopied] = useState(false);
+
+  const copyText = prepareCopyText(codeString);
+
+  const copy = () => {
+    trackCopyClicks(copyText);
+    setCopied(true);
+    setTimeout(() => {
+      setCopied(false);
+    }, 2000);
+  };
+  return (
+    <CopyToClipboard text={copyText} onCopy={copy}>
+      <Button
+        size="small"
+        variation="link"
+        disabled={copied}
+        className="code-copy"
+        testId={testId}
+        aria-describedby={title ? undefined : codeId}
+      >
+        {copied ? 'Copied!' : 'Copy'}
+        <VisuallyHidden>
+          {` `}
+          {title} code example
+        </VisuallyHidden>
+      </Button>
+    </CopyToClipboard>
+  );
+};

src/components/MDXComponents/__tests__/MDXCode.test.tsx

@@ -2,6 +2,12 @@ import * as React from 'react';
 import { render, screen } from '@testing-library/react';
 import { MDXCode } from '../MDXCode';
 
+jest.mock('react-copy-to-clipboard', () => ({
+  CopyToClipboard: jest.fn().mockImplementation(({ children }) => {
+    return children;
+  })
+}));
+
 const codeString = `
 import * as sns from 'aws-cdk-lib/aws-sns';
 import * as sqs from 'aws-cdk-lib/aws-sqs';

src/components/MDXComponents/__tests__/MDXCopyCodeButton.test.tsx

@@ -0,0 +1,93 @@
+import * as React from 'react';
+import { render, screen, waitFor } from '@testing-library/react';
+import { MDXCopyCodeButton, prepareCopyText } from '../MDXCopyCodeButton';
+import userEvent from '@testing-library/user-event';
+import * as trackModule from '../../../utils/track';
+
+const codeString = `
+import * as sns from 'aws-cdk-lib/aws-sns';
+
+// highlight-next-line
+import * as sqs from 'aws-cdk-lib/aws-sqs';
+import { defineBackend } from '@aws-amplify/backend';
+import { auth } from './auth/resource.js';
+import { data } from './data/resource.js';
+
+// highlight-start
+const backend = defineBackend({
+  auth,
+  data
+});
+// highlight-end
+
+const customResourceStack = backend.createStack('MyCustomResources');
+
+new sqs.Queue(customResourceStack, 'CustomQueue');
+new sns.Topic(customResourceStack, 'CustomTopic');
+`;
+
+const title = 'src/app.tsx';
+const codeId = 'codeId';
+const testId = 'copyButton';
+
+jest.mock('../../../utils/track', () => ({
+  trackCopyClicks: jest.fn().mockImplementation(() => codeString)
+}));
+
+jest.mock('copy-to-clipboard', () => {
+  return jest.fn();
+});
+
+describe('MDXCopyCodeButton', () => {
+  it('should render MDXCopyCodeButton', async () => {
+    render(
+      <MDXCopyCodeButton
+        codeString={codeString}
+        codeId={codeId}
+        testId={testId}
+        title={title}
+      />
+    );
+    const copyButton = await screen.findByTestId(testId);
+    expect(copyButton).toBeInTheDocument();
+  });
+
+  it('should track CopyCodeButton on click', async () => {
+    jest.spyOn(trackModule, 'trackCopyClicks');
+    render(
+      <MDXCopyCodeButton
+        codeString={codeString}
+        codeId={codeId}
+        testId={testId}
+        title={title}
+      />
+    );
+    const copyButton = await screen.findByTestId(testId);
+    userEvent.click(copyButton);
+
+    await waitFor(() => {
+      expect(trackModule.trackCopyClicks).toHaveBeenCalled();
+    });
+  });
+
+  it('should use aria-describedBy if no title is supplied', async () => {
+    render(
+      <MDXCopyCodeButton
+        codeString={codeString}
+        codeId={codeId}
+        testId={testId}
+      />
+    );
+    const copyButton = await screen.findByTestId(testId);
+    expect(copyButton).toHaveAttribute('aria-describedby', codeId);
+  });
+});
+
+describe('prepareCopyText', () => {
+  it('should return code string without markdown comments', () => {
+    const copyText = prepareCopyText(codeString);
+    expect(copyText.indexOf('// highlight-next-line')).toEqual(-1);
+    expect(copyText.indexOf('// highlight-start')).toEqual(-1);
+    expect(copyText.indexOf('// highlight-end')).toEqual(-1);
+  });
+});