From b9df3ba9548789b48ad25bd98666077d5c3b6ab8 Mon Sep 17 00:00:00 2001
From: "Ricardo Q. Bazan" <raquirozbazan@gmail.com>
Date: Mon, 6 Dec 2021 19:32:39 -0500
Subject: [PATCH] feat: add `flagQueryFn` to the context value (#10)

---
 src/index.tsx | 66 ++++++++++++++++++++++++++++-----------------------
 1 file changed, 36 insertions(+), 30 deletions(-)

diff --git a/src/index.tsx b/src/index.tsx
index 14db10e..035a677 100644
--- a/src/index.tsx
+++ b/src/index.tsx
@@ -18,6 +18,7 @@ export type FeatureProviderProps =
 
 export interface FeatureContextValue {
   cache: TCache
+  flagQueryFn: ReturnType<typeof createFlagQueryFn>
 }
 
 export const Op = {
@@ -50,34 +51,6 @@ export function createCache<F extends DefaultFeature = DefaultFeature>(features:
   }, {} as Record<string, F>)
 }
 
-export function FeatureProvider(props: FeatureProviderProps) {
-  // @ts-expect-error
-  const { features, cache, children } = props
-
-  const contextValue = React.useMemo(() => {
-    return { cache: cache ?? createCache(features) }
-  }, [cache, features])
-
-  return <FeatureContext.Provider value={contextValue}>{children}</FeatureContext.Provider>
-}
-
-function useToggledContext() {
-  const contextValue = React.useContext(FeatureContext)
-
-  // @ts-ignore
-  if (contextValue === NO_PROVIDER) {
-    throw new Error('Component must be wrapped with FeatureProvider.')
-  }
-
-  return contextValue
-}
-
-export function useFeature(slug: string) {
-  const { cache } = useToggledContext()
-
-  return cache[slug]
-}
-
 export function createFlagQueryFn(cache: Record<string, DefaultFeature>) {
   return function flagQueryFn(flagQuery: FlagQuery) {
     if (typeof flagQuery === 'string') {
@@ -131,10 +104,43 @@ export function createFlagQueryFn(cache: Record<string, DefaultFeature>) {
   }
 }
 
-export function useFlagQueryFn() {
+export function FeatureProvider(props: FeatureProviderProps) {
+  // @ts-expect-error
+  const { features, cache, children } = props
+
+  const contextValue = React.useMemo(() => {
+    const _cache = cache ?? createCache(features)
+
+    return {
+      cache: _cache,
+      flagQueryFn: createFlagQueryFn(_cache),
+    }
+  }, [cache, features])
+
+  return <FeatureContext.Provider value={contextValue}>{children}</FeatureContext.Provider>
+}
+
+function useToggledContext() {
+  const contextValue = React.useContext(FeatureContext)
+
+  // @ts-ignore
+  if (contextValue === NO_PROVIDER) {
+    throw new Error('Component must be wrapped with FeatureProvider.')
+  }
+
+  return contextValue
+}
+
+export function useFeature(slug: string) {
   const { cache } = useToggledContext()
 
-  return React.useMemo(() => createFlagQueryFn(cache), [cache])
+  return cache[slug]
+}
+
+export function useFlagQueryFn() {
+  const { flagQueryFn } = useToggledContext()
+
+  return flagQueryFn
 }
 
 export function useFlag(flagQuery: FlagQuery) {