diff --git a/packages/core/src/api/index.ts b/packages/core/src/api/index.ts
index e1bca9814e..4c77e94984 100644
--- a/packages/core/src/api/index.ts
+++ b/packages/core/src/api/index.ts
@@ -16,5 +16,9 @@
 
 export * from './api';
 export * from './apis';
-export { FeatureFlags, FeatureFlagsContext } from './app/FeatureFlags';
+export {
+  FeatureFlags,
+  FeatureFlagsContext,
+  FeatureFlagsContextProvider,
+} from './app/FeatureFlags';
 export { useApp } from './app/AppContext';
diff --git a/plugins/welcome/src/components/WelcomePage/ToggleFeatureFlagButton.test.tsx b/plugins/welcome/src/components/WelcomePage/ToggleFeatureFlagButton.test.tsx
index 8efb25f0a6..f9be280f50 100644
--- a/plugins/welcome/src/components/WelcomePage/ToggleFeatureFlagButton.test.tsx
+++ b/plugins/welcome/src/components/WelcomePage/ToggleFeatureFlagButton.test.tsx
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import React from 'react';
+import React, { ReactNode } from 'react';
 import { render, fireEvent } from '@testing-library/react';
 import ToggleFeatureFlagButton from './ToggleFeatureFlagButton';
 import {
@@ -22,8 +22,29 @@ import {
   featureFlagsApiRef,
   ApiProvider,
   FeatureFlags,
+  FeatureFlagsContextProvider,
 } from '@backstage/core';
 
+function withFeatureFlags(children: ReactNode) {
+  const featureFlags = new Set([
+    { pluginId: 'welcome', name: 'enable-welcome-box' },
+  ]);
+
+  return (
+    <FeatureFlagsContextProvider featureFlags={featureFlags}>
+      {children}
+    </FeatureFlagsContextProvider>
+  );
+}
+
+function withApiRegistry(children: ReactNode) {
+  return (
+    <ApiProvider apis={ApiRegistry.from([[featureFlagsApiRef, FeatureFlags]])}>
+      {children}
+    </ApiProvider>
+  );
+}
+
 describe('ToggleFeatureFlagButton', () => {
   beforeEach(() => {
     window.localStorage.clear();
@@ -31,11 +52,7 @@ describe('ToggleFeatureFlagButton', () => {
 
   it('should enable the feature flag', () => {
     const rendered = render(
-      <ApiProvider
-        apis={ApiRegistry.from([[featureFlagsApiRef, FeatureFlags]])}
-      >
-        <ToggleFeatureFlagButton />
-      </ApiProvider>,
+      withFeatureFlags(withApiRegistry(<ToggleFeatureFlagButton />)),
     );
 
     const button = rendered.getByTestId('button-switch-feature-flag-state');
@@ -50,11 +67,7 @@ describe('ToggleFeatureFlagButton', () => {
 
   it('should disable the feature flag', () => {
     const rendered = render(
-      <ApiProvider
-        apis={ApiRegistry.from([[featureFlagsApiRef, FeatureFlags]])}
-      >
-        <ToggleFeatureFlagButton />
-      </ApiProvider>,
+      withFeatureFlags(withApiRegistry(<ToggleFeatureFlagButton />)),
     );
 
     const button = rendered.getByTestId('button-switch-feature-flag-state');