diff --git a/src/resolvers/MutationResolvers.ts b/src/resolvers/MutationResolvers.ts
index adcc20c..9b48a7d 100644
--- a/src/resolvers/MutationResolvers.ts
+++ b/src/resolvers/MutationResolvers.ts
@@ -30,7 +30,7 @@ export const MutationResolvers: Resolvers<ServerContext> = {
           : ETANotificationScheduler.defaultSecondsThresholdForNotificationToFire,
       }
 
-      await context.notificationRepository.scheduleNotification(notificationData);
+      await context.notificationRepository.addOrUpdateNotification(notificationData);
 
       const response: NotificationResponse = {
         message: "Notification scheduled",
@@ -40,8 +40,9 @@ export const MutationResolvers: Resolvers<ServerContext> = {
       return response;
     },
     cancelNotification: async (_parent, args, context, _info) => {
-      if (context.notificationRepository.isNotificationScheduled(args.input)) {
-        await context.notificationRepository.cancelNotificationIfExists(args.input);
+      const isScheduled = await context.notificationRepository.isNotificationScheduled(args.input)
+      if (isScheduled) {
+        await context.notificationRepository.deleteNotificationIfExists(args.input);
         return {
           success: true,
           message: "Notification cancelled",
diff --git a/src/resolvers/QueryResolvers.ts b/src/resolvers/QueryResolvers.ts
index 996f9ac..06df1a9 100644
--- a/src/resolvers/QueryResolvers.ts
+++ b/src/resolvers/QueryResolvers.ts
@@ -18,11 +18,11 @@ export const QueryResolvers: Resolvers<ServerContext> = {
     },
     isNotificationScheduled: async (_parent, args, contextValue, _info) => {
       const notificationData = args.input;
-      return contextValue.notificationRepository.isNotificationScheduled(notificationData);
+      return await contextValue.notificationRepository.isNotificationScheduled(notificationData);
     },
     secondsThresholdForNotification: async (_parent, args, contextValue, _info) => {
       const notificationData = args.input;
-      return contextValue.notificationRepository.getSecondsThresholdForScheduledNotification(notificationData);
+      return await contextValue.notificationRepository.getSecondsThresholdForNotificationIfExists(notificationData);
     },
   },
 }
diff --git a/test/resolvers/MutationResolverTests.test.ts b/test/resolvers/MutationResolverTests.test.ts
index 0fe487a..055195e 100644
--- a/test/resolvers/MutationResolverTests.test.ts
+++ b/test/resolvers/MutationResolverTests.test.ts
@@ -38,13 +38,13 @@ describe("MutationResolvers", () => {
     }
     `
 
-    function assertFailedResponse(response: any, notificationInput: NotificationInput) {
+    async function assertFailedResponse(response: any, notificationInput: NotificationInput) {
       assert(response.body.kind === "single");
       expect(response.body.singleResult.errors).toBeUndefined();
       const notificationResponse = response.body.singleResult.data?.scheduleNotification as any;
       expect(notificationResponse.success).toBe(false);
 
-      expect(context.notificationRepository.isNotificationScheduled(notificationInput)).toBe(false);
+      expect(await context.notificationRepository.isNotificationScheduled(notificationInput)).toBe(false);
     }
 
 
@@ -72,7 +72,7 @@ describe("MutationResolvers", () => {
       expect(notificationResponse?.success).toBe(true);
       expect(notificationResponse?.data).toEqual(expectedNotificationData);
 
-      expect(context.notificationRepository.getSecondsThresholdForScheduledNotification(expectedNotificationData)).toBe(240);
+      expect(await context.notificationRepository.getSecondsThresholdForNotificationIfExists(expectedNotificationData)).toBe(240);
     });
 
     it("adds a notification with the default seconds threshold if none is provided", async () => {
@@ -93,7 +93,7 @@ describe("MutationResolvers", () => {
       const notificationResponse = response.body.singleResult.data?.scheduleNotification as any;
       expect(notificationResponse?.success).toBe(true);
 
-      expect(context.notificationRepository.getSecondsThresholdForScheduledNotification(notificationInput)).toBe(180);
+      expect(await context.notificationRepository.getSecondsThresholdForNotificationIfExists(notificationInput)).toBe(180);
     });
 
     it("fails if the shuttle ID doesn't exist", async () => {
@@ -106,7 +106,7 @@ describe("MutationResolvers", () => {
         stopId: stop.id,
       }
       const response = await getServerResponse(query, notificationInput);
-      assertFailedResponse(response, notificationInput);
+      await assertFailedResponse(response, notificationInput);
     });
 
     it("fails if the stop ID doesn't exist", async () => {
@@ -120,7 +120,7 @@ describe("MutationResolvers", () => {
       }
       const response = await getServerResponse(query, notificationInput);
 
-      assertFailedResponse(response, notificationInput);
+      await assertFailedResponse(response, notificationInput);
     });
   });
 
@@ -150,7 +150,7 @@ describe("MutationResolvers", () => {
         stopId: stop.id,
         secondsThreshold: 180,
       }
-      await context.notificationRepository.scheduleNotification(notificationInput);
+      await context.notificationRepository.addOrUpdateNotification(notificationInput);
 
       const notificationLookup = {
         ...notificationInput
@@ -166,7 +166,7 @@ describe("MutationResolvers", () => {
       expect(notificationResponse.success).toBe(true);
       expect(notificationResponse.data).toEqual(notificationLookup);
 
-      expect(context.notificationRepository.isNotificationScheduled(notificationLookup)).toBe(false);
+      expect(await context.notificationRepository.isNotificationScheduled(notificationLookup)).toBe(false);
     });
 
     it("fails if the notification doesn't exist", async () => {
diff --git a/test/resolvers/QueryResolverTests.test.ts b/test/resolvers/QueryResolverTests.test.ts
index 11a3282..52d6605 100644
--- a/test/resolvers/QueryResolverTests.test.ts
+++ b/test/resolvers/QueryResolverTests.test.ts
@@ -2,8 +2,8 @@ import { describe, expect, it } from "@jest/globals";
 import { generateMockSystems } from "../testHelpers/mockDataGenerators";
 import { setupTestServerContext, setupTestServerHolder } from "../testHelpers/apolloTestServerHelpers";
 import assert = require("node:assert");
-import { ScheduledNotification } from "../../src/notifications/schedulers/ETANotificationScheduler";
 import { addMockShuttleToRepository, addMockStopToRepository } from "../testHelpers/repositorySetupHelpers";
+import { ScheduledNotification } from "../../src/repositories/NotificationRepository";
 
 // See Apollo documentation for integration test guide
 // https://www.apollographql.com/docs/apollo-server/testing/testing
@@ -106,13 +106,13 @@ describe("QueryResolvers", () => {
       const shuttle = await addMockShuttleToRepository(context.shuttleRepository, "1");
       const stop = await addMockStopToRepository(context.shuttleRepository, "1")
 
-      const notification: NotificationSchedulingArguments = {
+      const notification: ScheduledNotification = {
         shuttleId: shuttle.id,
         stopId: stop.id,
         deviceId: "1",
         secondsThreshold: 240,
       };
-      await context.notificationRepository.scheduleNotification(notification);
+      await context.notificationRepository.addOrUpdateNotification(notification);
 
       const notificationLookup: any = {
         ...notification,