Mock hooks directly instead of context value

packages/keychain/.storybook/preview.tsx

@@ -1,10 +1,14 @@
 import React from "react";
 import type { Preview, ReactRenderer } from "@storybook/react";
 import { withThemeByClassName } from "@storybook/addon-themes";
-import { SonnerToaster } from "@cartridge/ui-next";
-
-import { StoryParameters } from "./mock";
-import { Provider } from "./provider";
+import { SonnerToaster, useThemeEffect } from "@cartridge/ui-next";
+import { BrowserRouter } from "react-router-dom";
+import { voyager, jsonRpcProvider, StarknetConfig } from "@starknet-react/core";
+import { sepolia, mainnet } from "@starknet-react/chains";
+import { num } from "starknet";
+import { createMockConnection } from "#hooks/connection.mock";
+import { controllerConfigs, defaultTheme } from "@cartridge/presets";
+import { UIProvider } from "#components/provider/ui";
 
 import "../src/index.css";
 
@@ -41,12 +45,38 @@ const preview: Preview = {
       },
       defaultTheme: "dark",
     }),
-    (Story, { parameters }) => (
-      <Provider parameters={parameters as StoryParameters}>
-        <Story />
-        <SonnerToaster />
-      </Provider>
-    ),
+    (Story, { parameters }) => {
+      {
+        /* TODO: Replace with mock once #1436 is merged */
+      }
+      useThemeEffect({
+        theme: parameters.preset
+          ? (controllerConfigs[parameters.preset].theme ?? defaultTheme)
+          : defaultTheme,
+        assetUrl: "",
+      });
+      const connection = createMockConnection();
+
+      return (
+        // Render only third party providers which consumer hook is directly used across the source code and hard to mock otherwise
+        <BrowserRouter>
+          <StarknetConfig
+            explorer={voyager}
+            chains={[sepolia, mainnet]}
+            defaultChainId={num.toBigInt(connection.controller!.chainId())}
+            provider={jsonRpcProvider({
+              rpc: () => ({ nodeUrl: connection.rpcUrl }),
+            })}
+          >
+            {/* TODO: Remove once #1436 is merged */}
+            <UIProvider>
+              <Story />
+            </UIProvider>
+            <SonnerToaster />
+          </StarknetConfig>
+        </BrowserRouter>
+      );
+    },
   ],
 };
 

packages/keychain/package.json

@@ -21,6 +21,18 @@
     "test:storybook": "start-server-and-test 'pnpm storybook:serve' 6001 'pnpm test-storybook -u'"
   },
   "imports": {
+    "#hooks/connection": {
+      "storybook": "./src/hooks/connection.mock.ts",
+      "default": "./src/hooks/connection.ts"
+    },
+    "#components/provider/posthog": {
+      "storybook": "./src/components/provider/posthog.mock.tsx",
+      "default": "./src/components/provider/posthog.tsx"
+    },
+    "#components/provider/upgrade": {
+      "storybook": "./src/components/provider/upgrade.mock.tsx",
+      "default": "./src/components/provider/upgrade.tsx"
+    },
     "#*": [
       "./src/*",
       "./src/*.ts",

packages/keychain/src/components/DeployController.stories.tsx

@@ -3,21 +3,23 @@ import type { Meta, StoryObj } from "@storybook/react";
 import { DeployController } from "./DeployController";
 import { constants, num } from "starknet";
 import { JsControllerError } from "@cartridge/account-wasm/controller";
-import Controller from "#utils/controller";
+import { useConnection, createMockConnection } from "#hooks/connection.mock";
 
 const meta = {
   component: DeployController,
-  parameters: {
-    connection: {
-      controller: {
-        chainId: () => constants.StarknetChainId.SN_SEPOLIA as string,
-        callContract: () =>
-          Promise.resolve([num.toHex("2000000000000000000"), "0x0"]),
-        username: () => "test-account",
-        address: () =>
-          "0x0000000000000000000000000000000000000000000000000000000000000000",
-      } as Partial<Controller>,
-    },
+  beforeEach: () => {
+    useConnection.mockReturnValue(
+      createMockConnection({
+        controller: {
+          chainId: () => constants.StarknetChainId.SN_SEPOLIA as string,
+          callContract: () =>
+            Promise.resolve([num.toHex("2000000000000000000"), "0x0"]),
+          username: () => "test-account",
+          address: () =>
+            "0x0000000000000000000000000000000000000000000000000000000000000000",
+        },
+      }),
+    );
   },
 } satisfies Meta<typeof DeployController>;
 

packages/keychain/src/components/ErrorBoundary.tsx

@@ -12,7 +12,7 @@ import { useConnection } from "#hooks/connection";
 import { CARTRIDGE_DISCORD_LINK } from "#const";
 import { Link } from "react-router-dom";
 import { useEffect } from "react";
-import { usePostHog } from "./provider/posthog";
+import { usePostHog } from "#components/provider/posthog";
 
 export class ErrorBoundary extends React.Component<
   PropsWithChildren,

packages/keychain/src/components/connect/CreateSession.stories.tsx

@@ -6,13 +6,6 @@ import { controllerConfigs } from "@cartridge/presets";
 
 const meta: Meta<typeof CreateSession> = {
   component: CreateSession,
-  parameters: {
-    connection: {
-      upgrade: {
-        isSynced: true,
-      },
-    },
-  },
 };
 
 export default meta;

packages/keychain/src/components/connect/CreateSession.tsx

@@ -4,7 +4,7 @@ import { Upgrade } from "#components/connect/Upgrade";
 import { UnverifiedSessionSummary } from "#components/session/UnverifiedSessionSummary";
 import { VerifiedSessionSummary } from "#components/session/VerifiedSessionSummary";
 import { NOW } from "#const";
-import { CreateSessionProvider } from "#context/session";
+import { CreateSessionProvider } from "#components/provider/session";
 import { useConnection } from "#hooks/connection";
 import {
   type ContractType,
@@ -24,7 +24,7 @@ import {
 import { useCallback, useMemo, useState } from "react";
 import { type BigNumberish } from "starknet";
 import { OcclusionDetector } from "#components/OcclusionDetector";
-import { useUpgrade } from "../provider/upgrade";
+import { useUpgrade } from "#components/provider/upgrade";
 
 const requiredPolicies: Array<ContractType> = ["VRF"];
 

packages/keychain/src/components/connect/RegisterSession.tsx

@@ -3,7 +3,7 @@ import { SessionConsent } from "#components/connect";
 import { UnverifiedSessionSummary } from "#components/session/UnverifiedSessionSummary";
 import { VerifiedSessionSummary } from "#components/session/VerifiedSessionSummary";
 import { NOW } from "#const";
-import { CreateSessionProvider } from "#context/session";
+import { CreateSessionProvider } from "#components/provider/session";
 import { useConnection } from "#hooks/connection";
 import {
   type ContractType,

packages/keychain/src/components/connect/SessionConsent.stories.tsx

@@ -1,14 +1,17 @@
 import type { Meta, StoryObj } from "@storybook/react";
 
 import { SessionConsent } from "./SessionConsent";
+import { useConnection, createMockConnection } from "#hooks/connection.mock";
 
 const meta = {
   // tags: ["autodocs"],
   component: SessionConsent,
-  parameters: {
-    connection: {
-      origin: "https://cartridge.gg",
-    },
+  beforeEach: () => {
+    useConnection.mockReturnValue(
+      createMockConnection({
+        origin: "https://cartridge.gg",
+      }),
+    );
   },
   args: {
     isVerified: true,

packages/keychain/src/components/connect/Upgrade.stories.tsx

@@ -1,23 +1,29 @@
 import type { Meta, StoryObj } from "@storybook/react";
-
 import { Upgrade } from "./Upgrade";
-import Controller from "#utils/controller";
-import { constants } from "starknet";
+import {
+  useUpgrade,
+  createMockUpgrade,
+} from "#components/provider/upgrade.mock";
+import { createMockConnection } from "#hooks/connection.mock";
+import { useConnection } from "#hooks/connection.mock";
 
 const meta = {
   component: Upgrade,
-  parameters: {
-    connection: {
-      controller: {
-        username: () => "Account 1",
-        chainId: () => constants.StarknetChainId.SN_MAIN,
-      } as Partial<Controller>,
-      upgrade: {
+  beforeEach: () => {
+    useUpgrade.mockReturnValue(
+      createMockUpgrade({
         latest: {
           changes: ["Update 1", "Update 2", "Update 3"],
         },
-      },
-    },
+      }),
+    );
+    useConnection.mockReturnValue(
+      createMockConnection({
+        controller: {
+          username: () => "Account 1",
+        },
+      }),
+    );
   },
 } satisfies Meta<typeof Upgrade>;
 

packages/keychain/src/components/connect/Upgrade.tsx

@@ -1,7 +1,7 @@
 import { LayoutContent, BoltIcon, CircleIcon } from "@cartridge/ui-next";
 import { ExecutionContainer } from "#components/ExecutionContainer";
 import { useConnection } from "#hooks/connection";
-import { useUpgrade } from "../provider/upgrade";
+import { useUpgrade } from "#components/provider/upgrade";
 
 export const Upgrade = () => {
   const { controller, logout } = useConnection();

packages/keychain/src/components/funding/index.stories.tsx

@@ -2,16 +2,19 @@ import type { Meta, StoryObj } from "@storybook/react";
 
 import { Funding } from ".";
 import { num } from "starknet";
+import { useConnection, createMockConnection } from "#hooks/connection.mock";
 
 const meta = {
   component: Funding,
-  parameters: {
-    connection: {
-      controller: {
-        callContract: () =>
-          Promise.resolve([num.toHex("2000000000000000000"), "0x0"]),
-      },
-    },
+  beforeEach: () => {
+    useConnection.mockReturnValue(
+      createMockConnection({
+        controller: {
+          callContract: () =>
+            Promise.resolve([num.toHex("2000000000000000000"), "0x0"]),
+        },
+      }),
+    );
   },
 } satisfies Meta<typeof Funding>;
 

packages/keychain/src/components/home.tsx

@@ -12,8 +12,8 @@ import { Settings } from "./settings";
 import { SignMessage } from "./SignMessage";
 import { PageLoading } from "./Loading";
 import { execute } from "#utils/connection/execute";
-import { useUpgrade } from "./provider/upgrade";
-import { usePostHog } from "./provider/posthog";
+import { usePostHog } from "#components/provider/posthog";
+import { useUpgrade } from "#components/provider/upgrade";
 
 export function Home() {
   const { context, setContext, controller, policies, origin } = useConnection();