diff --git a/src/database.ts b/src/database.ts
index f6e7595..3e2b948 100644
--- a/src/database.ts
+++ b/src/database.ts
@@ -348,23 +348,17 @@ export async function createClass(options: CreateClassOptions): Promise<CreateCl
           class_id: cls.id,
         });
 
-        classSetupRegistry.setupFunctions(storyName)?.forEach(setup => {
-          console.log(setup);
-          setup(cls, storyName);
-        });
+        const setupFunctions = classSetupRegistry.setupFunctions(storyName);
+        if (setupFunctions) {
+          for (const setupFunc of setupFunctions) {
+            await setupFunc(cls, storyName);
+          }
+        }
       }
 
       return cls;
     });
 
-    // Another piece of Hubble-specific functionality
-    // Note that we need to reload the class so that the virtual `small_class`
-    // column has its value populated
-    await cls.reload();
-    if (cls.asynchronous || cls.small_class) {
-      await addClassToMergeGroup(cls.id);
-    }
-
     return { result: result, class: creationInfo };
   } catch (error) {
     result = (error instanceof BaseError) ? createClassResultFromError(error) : CreateClassResult.Error;
@@ -941,19 +935,3 @@ export async function getDashboardGroupClasses(code: string): Promise<Class[] |
   });
 
 }
-
-export async function hubbleClassSetup(
-  cls: Class,
-  storyName: string,
-) {
-  if (cls) {
-    await ClassStories.create({
-      story_name: storyName,
-      class_id: cls.id
-    });
-
-    if (cls.asynchronous || cls.small_class) {
-      await addClassToMergeGroup(cls.id);
-    }
-  }
-}
diff --git a/src/registries.ts b/src/registries.ts
index 77e1021..abe42c8 100644
--- a/src/registries.ts
+++ b/src/registries.ts
@@ -1,6 +1,6 @@
 import { Class } from "./models";
 
-type ClassSetupFunction = (cls: Class, storyName: string) => void;
+type ClassSetupFunction = (cls: Class, storyName: string) => Promise<void>;
 
 class ClassSetupRegistry {