From 5f7bd0b399073b2f22b8eaec548c37ff19a5ea86 Mon Sep 17 00:00:00 2001
From: Jon Rohan <yes@jonrohan.codes>
Date: Fri, 4 Oct 2024 13:08:13 -0700
Subject: [PATCH] bugFix: When the feature flag is enabled and sx prop is
 passed in use, Box (#5068)

* When the feature flag is enabled and sx prop is passed in use, Box

* Create plenty-books-agree.md

* Adding tests

* Add as button back
---
 .changeset/plenty-books-agree.md              |  5 ++
 .../__tests__/toggleStyledComponent.test.tsx  | 49 +++++++++++++++++++
 .../internal/utils/toggleStyledComponent.tsx  | 13 ++++-
 3 files changed, 65 insertions(+), 2 deletions(-)
 create mode 100644 .changeset/plenty-books-agree.md
 create mode 100644 packages/react/src/internal/utils/__tests__/toggleStyledComponent.test.tsx

diff --git a/.changeset/plenty-books-agree.md b/.changeset/plenty-books-agree.md
new file mode 100644
index 00000000000..e893025c849
--- /dev/null
+++ b/.changeset/plenty-books-agree.md
@@ -0,0 +1,5 @@
+---
+"@primer/react": patch
+---
+
+fix for `toggleStyledComponent` utility, When the feature flag is enabled and sx prop is passed in use, Box
diff --git a/packages/react/src/internal/utils/__tests__/toggleStyledComponent.test.tsx b/packages/react/src/internal/utils/__tests__/toggleStyledComponent.test.tsx
new file mode 100644
index 00000000000..6a4206003d0
--- /dev/null
+++ b/packages/react/src/internal/utils/__tests__/toggleStyledComponent.test.tsx
@@ -0,0 +1,49 @@
+import React from 'react'
+import {render} from '@testing-library/react'
+import {FeatureFlags} from '../../../FeatureFlags'
+import {toggleStyledComponent} from '../toggleStyledComponent'
+import styled from 'styled-components'
+
+describe('toggleStyledComponent', () => {
+  test('renders the `as` prop when flag is enabled', () => {
+    const TestComponent = toggleStyledComponent('testFeatureFlag', () => <div />)
+    const {container} = render(
+      <FeatureFlags flags={{testFeatureFlag: true}}>
+        <TestComponent as="button" />
+      </FeatureFlags>,
+    )
+    expect(container.firstChild).toBeInstanceOf(HTMLButtonElement)
+  })
+
+  test('renders a div as fallback when flag is enabled and no `as` prop is provided', () => {
+    const TestComponent = toggleStyledComponent('testFeatureFlag', () => <div />)
+    const {container} = render(
+      <FeatureFlags flags={{testFeatureFlag: true}}>
+        <TestComponent />
+      </FeatureFlags>,
+    )
+    expect(container.firstChild).toBeInstanceOf(HTMLDivElement)
+  })
+
+  test('renders Box with `as` if `sx` is provided and flag is enabled', () => {
+    const TestComponent = toggleStyledComponent('testFeatureFlag', () => styled.div``)
+    const {container} = render(
+      <FeatureFlags flags={{testFeatureFlag: true}}>
+        <TestComponent as="button" sx={{color: 'red'}} />
+      </FeatureFlags>,
+    )
+
+    expect(container.firstChild).toBeInstanceOf(HTMLButtonElement)
+    expect(container.firstChild).toHaveStyle('color: red')
+  })
+
+  test('renders styled component when flag is disabled', () => {
+    const StyledComponent = toggleStyledComponent('testFeatureFlag', styled.div.attrs({['data-styled']: true})``)
+    const {container} = render(
+      <FeatureFlags flags={{testFeatureFlag: false}}>
+        <StyledComponent />
+      </FeatureFlags>,
+    )
+    expect(container.firstChild).toHaveAttribute('data-styled')
+  })
+})
diff --git a/packages/react/src/internal/utils/toggleStyledComponent.tsx b/packages/react/src/internal/utils/toggleStyledComponent.tsx
index 4a9bbd0e00a..7b366cb0eb0 100644
--- a/packages/react/src/internal/utils/toggleStyledComponent.tsx
+++ b/packages/react/src/internal/utils/toggleStyledComponent.tsx
@@ -1,9 +1,12 @@
 import React from 'react'
 import {useFeatureFlag} from '../../FeatureFlags'
+import Box from '../../Box'
+import {defaultSxProp} from '../../utils/defaultSxProp'
 
 type CSSModulesProps = {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   as?: string | React.ComponentType<any>
+  sx?: React.CSSProperties
 }
 
 /**
@@ -18,12 +21,18 @@ type CSSModulesProps = {
  * is disabled
  */
 export function toggleStyledComponent<T, P extends CSSModulesProps>(flag: string, Component: React.ComponentType<P>) {
-  const Wrapper = React.forwardRef<T, P>(function Wrapper({as: BaseComponent = 'div', ...rest}, ref) {
+  const Wrapper = React.forwardRef<T, P>(function Wrapper(
+    {as: BaseComponent = 'div', sx: sxProp = defaultSxProp, ...rest},
+    ref,
+  ) {
     const enabled = useFeatureFlag(flag)
     if (enabled) {
+      if (sxProp !== defaultSxProp) {
+        return <Box as={BaseComponent} {...rest} sx={sxProp} ref={ref} />
+      }
       return <BaseComponent {...rest} ref={ref} />
     }
-    return <Component as={BaseComponent} {...(rest as P)} ref={ref} />
+    return <Component as={BaseComponent} {...(rest as P)} sx={sxProp} ref={ref} />
   })
 
   return Wrapper