From 82fea932283012cc21a6202712b3bf1b3b41b818 Mon Sep 17 00:00:00 2001
From: Jon Rohan <yes@jonrohan.codes>
Date: Thu, 17 Oct 2024 16:10:14 -0700
Subject: [PATCH] Update Button.test.tsx

---
 .../src/Button/__tests__/Button.test.tsx      | 83 ++++++++++---------
 1 file changed, 42 insertions(+), 41 deletions(-)

diff --git a/packages/react/src/Button/__tests__/Button.test.tsx b/packages/react/src/Button/__tests__/Button.test.tsx
index bea2fb83a6a..3a5a253ff0e 100644
--- a/packages/react/src/Button/__tests__/Button.test.tsx
+++ b/packages/react/src/Button/__tests__/Button.test.tsx
@@ -25,6 +25,48 @@ const StatefulLoadingButton = (props: StatefulLoadingButtonProps) => {
   return <Button loading={isLoading} onClick={handleClick} {...props} />
 }
 
+describe('IconButton', () => {
+  it('should support `className` on the outermost element', () => {
+    const Element = () => <IconButton className={'test-class-name'} icon={SearchIcon} aria-label="Search button" />
+    const FeatureFlagElement = () => {
+      return (
+        <FeatureFlags
+          flags={{
+            primer_react_css_modules_team: true,
+            primer_react_css_modules_staff: true,
+            primer_react_css_modules_ga: true,
+          }}
+        >
+          <Element />
+        </FeatureFlags>
+      )
+    }
+    expect(render(<Element />).container.firstChild).toHaveClass('test-class-name')
+    expect(render(<FeatureFlagElement />).container.firstChild).toHaveClass('test-class-name')
+  })
+})
+
+describe('LinkButton', () => {
+  it('should support `className` on the outermost element', () => {
+    const Element = () => <LinkButton className={'test-class-name'} />
+    const FeatureFlagElement = () => {
+      return (
+        <FeatureFlags
+          flags={{
+            primer_react_css_modules_team: true,
+            primer_react_css_modules_staff: true,
+            primer_react_css_modules_ga: true,
+          }}
+        >
+          <Element />
+        </FeatureFlags>
+      )
+    }
+    expect(render(<Element />).container.firstChild).toHaveClass('test-class-name')
+    expect(render(<FeatureFlagElement />).container.firstChild).toHaveClass('test-class-name')
+  })
+})
+
 describe('Button', () => {
   behavesAsComponent({
     Component: TestButton,
@@ -320,45 +362,4 @@ describe('Button', () => {
     expect(tooltipEl).toBeInTheDocument()
     expect(triggerEl.getAttribute('aria-describedby')).toEqual(expect.stringContaining(tooltipEl.id))
   })
-
-  describe('with primer_react_css_modules_ga enabled', () => {
-    it('iconbutton should support custom `className` along with default classnames', () => {
-      const {container} = render(
-        <FeatureFlags
-          flags={{
-            primer_react_css_modules_ga: true,
-          }}
-        >
-          <IconButton className="test" aria-label="Test" icon={HeartIcon} />
-        </FeatureFlags>,
-      )
-      expect(container.firstChild).toHaveClass('IconButton')
-    })
-
-    it('button should support custom `className` along with default classnames', () => {
-      const {container} = render(
-        <FeatureFlags
-          flags={{
-            primer_react_css_modules_ga: true,
-          }}
-        >
-          <Button className="test">Hello</Button>
-        </FeatureFlags>,
-      )
-      expect(container.firstChild).toHaveClass('ButtonBase')
-    })
-
-    it('linkbutton should support custom `className` along with default classnames', () => {
-      const {container} = render(
-        <FeatureFlags
-          flags={{
-            primer_react_css_modules_ga: true,
-          }}
-        >
-          <LinkButton className="test">Hello</LinkButton>
-        </FeatureFlags>,
-      )
-      expect(container.firstChild).toHaveClass('ButtonBase')
-    })
-  })
 })