Use app context in organizationLogic too (#6565)

* Use app context in `organizationLogic` too

* Add `organizationLogic.test.ts`

frontend/src/lib/api.mock.ts

@@ -50,6 +50,10 @@ export function defaultAPIMocks(
             ingested_event: true,
             completed_snippet_onboarding: true,
         }
+    } else if (pathname === 'api/organizations/@current') {
+        return {
+            id: 'ABCD', // Should be a UUID but that doesn't matter here
+        }
     } else if (
         [
             'api/action/',

frontend/src/lib/utils/getAppContext.ts

@@ -1,5 +1,11 @@
 import { AppContext, PathType } from '~/types'
 
+declare global {
+    export interface Window {
+        POSTHOG_APP_CONTEXT?: AppContext
+    }
+}
+
 export function getAppContext(): AppContext | undefined {
     return (window as any)['POSTHOG_APP_CONTEXT'] || undefined
 }

frontend/src/scenes/organizationLogic.test.ts

@@ -0,0 +1,56 @@
+import { BuiltLogic } from 'kea'
+import { expectLogic } from 'kea-test-utils'
+import { initKeaTestLogic } from '~/test/init'
+import { defaultAPIMocks, mockAPI } from '../lib/api.mock'
+import { AppContext } from '../types'
+import { organizationLogic, OrganizationUpdatePayload } from './organizationLogic'
+import { organizationLogicType } from './organizationLogicType'
+
+jest.mock('lib/api')
+
+describe('entityFilterLogic', () => {
+    let logic: BuiltLogic<organizationLogicType<OrganizationUpdatePayload>>
+
+    mockAPI(async (url) => {
+        return defaultAPIMocks(url)
+    })
+
+    describe('if POSTHOG_APP_CONTEXT available', () => {
+        beforeEach(() => {
+            window.POSTHOG_APP_CONTEXT = { current_user: { organization: { id: 'WXYZ' } } } as unknown as AppContext
+        })
+        afterEach(() => {
+            delete window.POSTHOG_APP_CONTEXT
+        })
+        initKeaTestLogic({
+            logic: organizationLogic,
+            onLogic: (l) => {
+                logic = l
+            },
+        })
+
+        it('loads organization from window', async () => {
+            await expectLogic(logic).toNotHaveDispatchedActions(['loadCurrentOrganization'])
+            await expectLogic(logic).toDispatchActions(['loadCurrentOrganizationSuccess'])
+            await expectLogic(logic).toMatchValues({
+                currentOrganization: { id: 'WXYZ' },
+            })
+        })
+    })
+
+    describe('if POSTHOG_APP_CONTEXT not available', () => {
+        initKeaTestLogic({
+            logic: organizationLogic,
+            onLogic: (l) => {
+                logic = l
+            },
+        })
+
+        it('loads organization from API', async () => {
+            await expectLogic(logic).toDispatchActions(['loadCurrentOrganization', 'loadCurrentOrganizationSuccess'])
+            await expectLogic(logic).toMatchValues({
+                currentOrganization: { id: 'ABCD' },
+            })
+        })
+    })
+})

frontend/src/scenes/organizationLogic.tsx

@@ -4,8 +4,9 @@ import { organizationLogicType } from './organizationLogicType'
 import { AvailableFeature, OrganizationType } from '~/types'
 import { toast } from 'react-toastify'
 import { userLogic } from './userLogic'
+import { getAppContext } from '../lib/utils/getAppContext'
 
-type OrganizationUpdatePayload = Partial<
+export type OrganizationUpdatePayload = Partial<
     Pick<OrganizationType, 'name' | 'personalization' | 'domain_whitelist' | 'is_member_join_email_enabled'>
 >
 
@@ -80,6 +81,16 @@ export const organizationLogic = kea<organizationLogicType<OrganizationUpdatePay
         },
     }),
     events: ({ actions }) => ({
-        afterMount: [actions.loadCurrentOrganization],
+        afterMount: () => {
+            const appContext = getAppContext()
+            const contextualOrganization = appContext?.current_user?.organization
+            if (contextualOrganization) {
+                // If app context is available (it should be practically always) we can immediately know currentOrganization
+                actions.loadCurrentOrganizationSuccess(contextualOrganization)
+            } else {
+                // If app context is not available, a traditional request is needed
+                actions.loadCurrentOrganization()
+            }
+        },
     }),
 })

frontend/src/scenes/teamLogic.tsx

@@ -117,9 +117,10 @@ export const teamLogic = kea<teamLogicType>({
     events: ({ actions }) => ({
         afterMount: () => {
             const appContext = getAppContext()
-            if (appContext) {
+            const contextualTeam = appContext?.current_team
+            if (contextualTeam) {
                 // If app context is available (it should be practically always) we can immediately know currentTeam
-                actions.loadCurrentTeamSuccess(appContext.current_team)
+                actions.loadCurrentTeamSuccess(contextualTeam)
             } else {
                 // If app context is not available, a traditional request is needed
                 actions.loadCurrentTeam()