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

@@ -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

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')
+  })
+})

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