Updated more tests

packages/design-system-react-native/src/hocs/withThemeProvider.test.tsx

@@ -1,19 +1,24 @@
 import {
   ThemeProvider,
-  useThemeContext,
+  ThemeContext,
   ColorSet,
   Theme,
 } from '@metamask/design-system-twrnc-preset';
 import { render } from '@testing-library/react-native';
-import React from 'react';
+import React, { forwardRef, useContext, createRef } from 'react';
 import { Text } from 'react-native';
 
 import { withThemeProvider } from './withThemeProvider';
 
-const TestThemeComponent = () => {
-  const { theme } = useThemeContext();
-  return <Text>{theme}</Text>;
-};
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+const TestThemeComponent = forwardRef((props, ref) => {
+  const themeContext = useContext(ThemeContext);
+  return (
+    <Text ref={ref as React.Ref<Text>}>
+      {themeContext.theme ? themeContext.theme : 'No Theme'}
+    </Text>
+  );
+});
 const WrappedComponent = withThemeProvider(TestThemeComponent);
 
 describe('withThemeProvider HOC', () => {
@@ -28,7 +33,12 @@ describe('withThemeProvider HOC', () => {
         <WrappedComponent />
       </ThemeProvider>,
     );
-
     expect(getByText(Theme.Dark)).toBeDefined();
   });
+
+  it('forwards ref to the wrapped component', () => {
+    const ref = createRef<Text>();
+    render(<WrappedComponent ref={ref} />);
+    expect(ref.current).toBeDefined();
+  });
 });

packages/design-system-react-native/src/hocs/withThemeProvider.tsx

@@ -1,37 +1,49 @@
-/* eslint-disable jsdoc/require-returns */
-// src/hocs/withThemeProvider.tsx
-
 import {
   ThemeProvider,
   ThemeContext,
   ColorSet,
   Theme,
 } from '@metamask/design-system-twrnc-preset';
-import React, { forwardRef, useContext } from 'react';
+import { render } from '@testing-library/react-native';
+import React from 'react';
+import { Text } from 'react-native';
+
+import { withThemeProvider } from './withThemeProvider';
 
-/**
- * HOC to wrap components with ThemeProvider if none is present.
- * @param Component - The component to wrap with ThemeProvider.
- */
-export function withThemeProvider<Props extends object>(
-  Component: React.ComponentType<Props>,
-) {
-  const WrappedComponent = forwardRef<unknown, Props>((props, ref) => {
-    // Check if a ThemeProvider is already present
-    const themeContext = useContext(ThemeContext);
+const TestThemeComponent = React.forwardRef((props, ref) => {
+  const themeContext = React.useContext(ThemeContext);
+  return (
+    <Text ref={ref}>
+      {themeContext.theme ? themeContext.theme : 'No Theme'}
+    </Text>
+  );
+});
+const WrappedComponent = withThemeProvider(TestThemeComponent);
 
-    // If ThemeProvider exists, use the component as is
-    if (themeContext) {
-      return <Component {...(props as Props)} ref={ref} />;
-    }
+describe('withThemeProvider HOC', () => {
+  it('provides default theme when no ThemeProvider is present', () => {
+    const { getByText } = render(<WrappedComponent />);
+    expect(getByText(Theme.Default)).toBeDefined();
+  });
 
-    // Otherwise, wrap with ThemeProvider
-    return (
-      <ThemeProvider colorSet={ColorSet.Brand} theme={Theme.Default}>
-        <Component {...(props as Props)} ref={ref} />
-      </ThemeProvider>
+  it('does not override existing theme context', () => {
+    const { getByText } = render(
+      <ThemeProvider colorSet={ColorSet.Brand} theme={Theme.Dark}>
+        <WrappedComponent />
+      </ThemeProvider>,
     );
+    expect(getByText(Theme.Dark)).toBeDefined();
   });
 
-  return WrappedComponent;
-}
+  it('forwards ref to the wrapped component', () => {
+    const ref = React.createRef();
+    render(<WrappedComponent ref={ref} />);
+    expect(ref.current).toBeDefined();
+  });
+
+  it('passes additional props to the wrapped component', () => {
+    const { getByText } = render(<WrappedComponent testProp="testValue" />);
+    // Update TestThemeComponent to display testProp for this test
+    expect(getByText('testValue')).toBeDefined();
+  });
+});

packages/design-system-twrnc-preset/src/Theme/Theme.hooks.ts

@@ -6,13 +6,3 @@ export const useTailwind = () => {
   const { tw } = useContext(ThemeContext);
   return tw;
 };
-
-export const useThemeContext = () => {
-  const context = useContext(ThemeContext);
-
-  if (!context) {
-    throw new Error('useThemeContext must be used within a ThemeProvider');
-  }
-
-  return context;
-};

packages/design-system-twrnc-preset/src/Theme/index.ts

@@ -1,4 +1,4 @@
-export { useTailwind, useThemeContext } from './Theme.hooks';
+export { useTailwind } from './Theme.hooks';
 export { ThemeContext, ThemeProvider } from './Theme.providers';
 export { ColorSet, Theme } from './Theme.types';
 export type { ThemeContextProps, ThemeProviderProps } from './Theme.types';

packages/design-system-twrnc-preset/src/index.ts

@@ -1,7 +1,6 @@
 export type { ThemeContextProps, ThemeProviderProps } from './Theme';
 export {
   useTailwind,
-  useThemeContext,
   ThemeContext,
   ThemeProvider,
   ColorSet,