From 03f9a4c5b345e00e3cafc4d3a1b813f85a3f7a09 Mon Sep 17 00:00:00 2001
From: Ben White <ben@posthog.com>
Date: Thu, 20 Jun 2024 17:52:33 +0200
Subject: [PATCH] Fixes

---
 plugin-server/src/cdp/hog-watcher.ts          | 16 +++-------
 .../tests/cdp/hog-function-manager.test.ts    | 32 +++++++++++--------
 2 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/plugin-server/src/cdp/hog-watcher.ts b/plugin-server/src/cdp/hog-watcher.ts
index 9036956520755..ea86442f34ab2 100644
--- a/plugin-server/src/cdp/hog-watcher.ts
+++ b/plugin-server/src/cdp/hog-watcher.ts
@@ -132,6 +132,8 @@ export class HogWatcherObserver {
             state,
         })
 
+        // TODO: Somehow report this back to PostHog so we can display it in the UI
+
         return state
     }
 
@@ -225,22 +227,12 @@ export class HogWatcher {
         })
     }
 
-    getOverflowedHogFunctionIds(): HogFunctionType['id'][] {
-        // TODO
-        return []
-    }
-
     isHogFunctionOverflowed(hogFunctionId: HogFunctionType['id']): boolean {
-        return this.getOverflowedHogFunctionIds().includes(hogFunctionId)
-    }
-
-    getDisabledHogFunctionIds(): HogFunctionType['id'][] {
-        // TODO
-        return []
+        return this.getObserver(hogFunctionId).currentState() === HogWatcherState.overflowed
     }
 
     isHogFunctionDisabled(hogFunctionId: HogFunctionType['id']): boolean {
-        return this.getDisabledHogFunctionIds().includes(hogFunctionId)
+        return this.getObserver(hogFunctionId).currentState() >= HogWatcherState.disabledForPeriod
     }
 
     private getObserver(id: HogFunctionType['id']): HogWatcherObserver {
diff --git a/plugin-server/tests/cdp/hog-function-manager.test.ts b/plugin-server/tests/cdp/hog-function-manager.test.ts
index a3245b4ef97c1..aae8ba81b2c67 100644
--- a/plugin-server/tests/cdp/hog-function-manager.test.ts
+++ b/plugin-server/tests/cdp/hog-function-manager.test.ts
@@ -47,10 +47,10 @@ describe('HogFunctionManager', () => {
     })
 
     it('returns the hog functions', async () => {
-        let functionsMap = manager.getTeamHogFunctions(teamId1)
+        let items = manager.getTeamHogFunctions(teamId1)
 
-        expect(functionsMap).toEqual({
-            [hogFunctions[0].id]: {
+        expect(items).toEqual([
+            {
                 id: hogFunctions[0].id,
                 team_id: teamId1,
                 name: 'Test Hog Function team 1',
@@ -59,7 +59,7 @@ describe('HogFunctionManager', () => {
                 bytecode: null,
                 filters: null,
             },
-        })
+        ])
 
         await hub.db.postgres.query(
             PostgresUse.COMMON_WRITE,
@@ -71,18 +71,24 @@ describe('HogFunctionManager', () => {
         // This is normally dispatched by django
         await manager.reloadHogFunctions(teamId1, [hogFunctions[0].id])
 
-        functionsMap = manager.getTeamHogFunctions(teamId1)
+        items = manager.getTeamHogFunctions(teamId1)
 
-        expect(functionsMap[hogFunctions[0].id]).toMatchObject({
-            id: hogFunctions[0].id,
-            name: 'Test Hog Function team 1 updated',
-        })
+        expect(items).toMatchObject([
+            {
+                id: hogFunctions[0].id,
+                name: 'Test Hog Function team 1 updated',
+            },
+        ])
     })
 
     it('removes disabled functions', async () => {
-        let functionsMap = manager.getTeamHogFunctions(teamId1)
+        let items = manager.getTeamHogFunctions(teamId1)
 
-        expect(functionsMap).toHaveProperty(hogFunctions[0].id)
+        expect(items).toMatchObject([
+            {
+                id: hogFunctions[0].id,
+            },
+        ])
 
         await hub.db.postgres.query(
             PostgresUse.COMMON_WRITE,
@@ -94,8 +100,8 @@ describe('HogFunctionManager', () => {
         // This is normally dispatched by django
         await manager.reloadHogFunctions(teamId1, [hogFunctions[0].id])
 
-        functionsMap = manager.getTeamHogFunctions(teamId1)
+        items = manager.getTeamHogFunctions(teamId1)
 
-        expect(functionsMap).not.toHaveProperty(hogFunctions[0].id)
+        expect(items).toEqual([])
     })
 })