diff --git a/dist/restore/index.js b/dist/restore/index.js
index 589163b..55d8cd2 100644
--- a/dist/restore/index.js
+++ b/dist/restore/index.js
@@ -62949,7 +62949,7 @@ module.exports.implForWrapper = function (wrapper) {
 "use strict";
 
 Object.defineProperty(exports, "__esModule", ({ value: true }));
-exports.State = exports.Outputs = exports.Inputs = void 0;
+exports.RefKey = exports.Gatsby = exports.Platform = exports.Events = exports.State = exports.Outputs = exports.Inputs = void 0;
 var Inputs;
 (function (Inputs) {
     Inputs["UseCache"] = "use-cache";
@@ -62965,6 +62965,21 @@ var State;
     State["CachePrimaryKey"] = "CACHE_KEY";
     State["CacheMatchedKey"] = "CACHE_RESULT";
 })(State = exports.State || (exports.State = {}));
+var Events;
+(function (Events) {
+    Events["Key"] = "GITHUB_EVENT_NAME";
+})(Events = exports.Events || (exports.Events = {}));
+var Platform;
+(function (Platform) {
+    Platform["RunnerOs"] = "RUNNER_OS";
+})(Platform = exports.Platform || (exports.Platform = {}));
+var Gatsby;
+(function (Gatsby) {
+    Gatsby["CacheDir"] = ".cache";
+    Gatsby["PublicDir"] = "public";
+    Gatsby["Env"] = "GATSBY_EXPERIMENTAL_PAGE_BUILD_ON_DATA_CHANGES";
+})(Gatsby = exports.Gatsby || (exports.Gatsby = {}));
+exports.RefKey = "GITHUB_REF";
 
 
 /***/ }),
@@ -62974,6 +62989,26 @@ var State;
 
 "use strict";
 
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", ({ value: true }));
+const restoreImpl_1 = __importDefault(__nccwpck_require__(2357));
+const stateProvider_1 = __nccwpck_require__(1527);
+async function run() {
+    await (0, restoreImpl_1.default)(new stateProvider_1.StateProvider());
+}
+void run();
+exports["default"] = run;
+
+
+/***/ }),
+
+/***/ 2357:
+/***/ (function(__unused_webpack_module, exports, __nccwpck_require__) {
+
+"use strict";
+
 var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
     if (k2 === undefined) k2 = k;
     var desc = Object.getOwnPropertyDescriptor(m, k);
@@ -63002,10 +63037,15 @@ const cache = __importStar(__nccwpck_require__(7799));
 const core = __importStar(__nccwpck_require__(2186));
 const utils = __importStar(__nccwpck_require__(6850));
 const constants_1 = __nccwpck_require__(9042);
-async function run() {
+async function restoreImpl(stateProvider) {
     try {
         if (!utils.isCacheFeatureAvailable()) {
-            utils.setCacheHitOutput(false);
+            core.setOutput(constants_1.Outputs.CacheHit, "false");
+            return;
+        }
+        if (!utils.isValidEvent()) {
+            const eventName = process.env[constants_1.Events.Key] || "";
+            utils.logWarning(`Event Validation Error: The event type ${eventName} is not supported because it's not tied to a branch or tag ref.`);
             return;
         }
         const useCache = core.getBooleanInput(constants_1.Inputs.UseCache);
@@ -63014,26 +63054,77 @@ async function run() {
         const restoreKeys = utils.getInputAsArray(constants_1.Inputs.RestoreKeys);
         let primaryKey = core.getInput(constants_1.Inputs.Key);
         if (!primaryKey) {
-            const platform = process.env.RUNNER_OS;
+            const platform = process.env[constants_1.Platform.RunnerOs];
             primaryKey = `${platform}-gatsby-build-`;
         }
         core.debug(`primary key is ${primaryKey}`);
-        core.saveState(constants_1.State.CachePrimaryKey, primaryKey);
+        stateProvider.setState(constants_1.State.CachePrimaryKey, primaryKey);
         const cacheKey = await cache.restoreCache(cachePaths, primaryKey, restoreKeys);
         if (!cacheKey) {
             core.info(`Cache not found for keys: ${[primaryKey, ...restoreKeys].join(", ")}`);
             return;
         }
-        utils.setCacheState(cacheKey);
+        stateProvider.setState(constants_1.State.CacheMatchedKey, cacheKey);
         const isExactKeyMatch = utils.isExactKeyMatch(primaryKey, cacheKey);
-        utils.setCacheHitOutput(isExactKeyMatch);
+        core.setOutput(constants_1.Outputs.CacheHit, isExactKeyMatch.toString());
         core.info(`Cache restored from key: ${cacheKey}`);
+        return cacheKey;
     }
     catch (error) {
         core.setFailed(error.message);
+        return;
     }
 }
-void run();
+exports["default"] = restoreImpl;
+
+
+/***/ }),
+
+/***/ 1527:
+/***/ (function(__unused_webpack_module, exports, __nccwpck_require__) {
+
+"use strict";
+
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
+    __setModuleDefault(result, mod);
+    return result;
+};
+Object.defineProperty(exports, "__esModule", ({ value: true }));
+exports.StateProvider = void 0;
+const core = __importStar(__nccwpck_require__(2186));
+const constants_1 = __nccwpck_require__(9042);
+class StateProvider {
+    getCacheState() {
+        const cacheKey = this.getState(constants_1.State.CacheMatchedKey);
+        if (cacheKey) {
+            core.debug(`Cache state/key: ${cacheKey}`);
+            return cacheKey;
+        }
+        return undefined;
+    }
+    getState = core.getState;
+    setState = core.saveState;
+}
+exports.StateProvider = StateProvider;
 
 
 /***/ }),
@@ -63070,7 +63161,7 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", ({ value: true }));
-exports.getBuildOutputPaths = exports.getInputAsArray = exports.isCacheFeatureAvailable = exports.logWarning = exports.getCacheState = exports.setCacheHitOutput = exports.setCacheState = exports.setBuildMode = exports.isExactKeyMatch = exports.isGhes = void 0;
+exports.getBuildOutputPaths = exports.setBuildMode = exports.isCacheFeatureAvailable = exports.getInputAsArray = exports.isValidEvent = exports.logWarning = exports.isExactKeyMatch = exports.isGhes = void 0;
 const cache = __importStar(__nccwpck_require__(7799));
 const core = __importStar(__nccwpck_require__(2186));
 const path_1 = __importDefault(__nccwpck_require__(1017));
@@ -63088,31 +63179,23 @@ function isExactKeyMatch(key, cacheKey) {
         }) === 0);
 }
 exports.isExactKeyMatch = isExactKeyMatch;
-function setBuildMode(useCache) {
-    process.env.GATSBY_EXPERIMENTAL_PAGE_BUILD_ON_DATA_CHANGES = String(useCache);
-}
-exports.setBuildMode = setBuildMode;
-function setCacheState(state) {
-    core.saveState(constants_1.State.CacheMatchedKey, state);
-}
-exports.setCacheState = setCacheState;
-function setCacheHitOutput(isCacheHit) {
-    core.setOutput(constants_1.Outputs.CacheHit, isCacheHit.toString());
-}
-exports.setCacheHitOutput = setCacheHitOutput;
-function getCacheState() {
-    const cacheKey = core.getState(constants_1.State.CacheMatchedKey);
-    if (cacheKey) {
-        core.debug(`Cache state/key: ${cacheKey}`);
-        return cacheKey;
-    }
-    return undefined;
-}
-exports.getCacheState = getCacheState;
 function logWarning(message) {
-    core.info(`[warning]${message}`);
+    const warningPrefix = "[warning]";
+    core.info(`${warningPrefix}${message}`);
 }
 exports.logWarning = logWarning;
+function isValidEvent() {
+    return constants_1.RefKey in process.env && Boolean(process.env[constants_1.RefKey]);
+}
+exports.isValidEvent = isValidEvent;
+function getInputAsArray(name, options) {
+    return core
+        .getInput(name, options)
+        .split("\n")
+        .map((s) => s.replace(/^!\s+/, "!").trim())
+        .filter((x) => x !== "");
+}
+exports.getInputAsArray = getInputAsArray;
 function isCacheFeatureAvailable() {
     if (cache.isFeatureAvailable()) {
         return true;
@@ -63126,16 +63209,12 @@ Otherwise please upgrade to GHES version >= 3.5 and If you are also using Github
     return false;
 }
 exports.isCacheFeatureAvailable = isCacheFeatureAvailable;
-function getInputAsArray(name, options) {
-    return core
-        .getInput(name, options)
-        .split("\n")
-        .map((s) => s.replace(/^!\s+/, "!").trim())
-        .filter((x) => x !== "");
+function setBuildMode(useCache) {
+    process.env[constants_1.Gatsby.Env] = String(useCache);
 }
-exports.getInputAsArray = getInputAsArray;
+exports.setBuildMode = setBuildMode;
 async function getBuildOutputPaths() {
-    const targetPaths = [".cache", "public"];
+    const targetPaths = [constants_1.Gatsby.CacheDir, constants_1.Gatsby.PublicDir];
     const buildOutputPaths = [];
     for await (const target of targetPaths) {
         buildOutputPaths.push(path_1.default.join(process.cwd(), target));
diff --git a/dist/save/index.js b/dist/save/index.js
index 6ba42c5..8c8a5b8 100644
--- a/dist/save/index.js
+++ b/dist/save/index.js
@@ -62949,7 +62949,7 @@ module.exports.implForWrapper = function (wrapper) {
 "use strict";
 
 Object.defineProperty(exports, "__esModule", ({ value: true }));
-exports.State = exports.Outputs = exports.Inputs = void 0;
+exports.RefKey = exports.Gatsby = exports.Platform = exports.Events = exports.State = exports.Outputs = exports.Inputs = void 0;
 var Inputs;
 (function (Inputs) {
     Inputs["UseCache"] = "use-cache";
@@ -62965,6 +62965,21 @@ var State;
     State["CachePrimaryKey"] = "CACHE_KEY";
     State["CacheMatchedKey"] = "CACHE_RESULT";
 })(State = exports.State || (exports.State = {}));
+var Events;
+(function (Events) {
+    Events["Key"] = "GITHUB_EVENT_NAME";
+})(Events = exports.Events || (exports.Events = {}));
+var Platform;
+(function (Platform) {
+    Platform["RunnerOs"] = "RUNNER_OS";
+})(Platform = exports.Platform || (exports.Platform = {}));
+var Gatsby;
+(function (Gatsby) {
+    Gatsby["CacheDir"] = ".cache";
+    Gatsby["PublicDir"] = "public";
+    Gatsby["Env"] = "GATSBY_EXPERIMENTAL_PAGE_BUILD_ON_DATA_CHANGES";
+})(Gatsby = exports.Gatsby || (exports.Gatsby = {}));
+exports.RefKey = "GITHUB_REF";
 
 
 /***/ }),
@@ -62974,6 +62989,26 @@ var State;
 
 "use strict";
 
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", ({ value: true }));
+const saveImpl_1 = __importDefault(__nccwpck_require__(6589));
+const stateProvider_1 = __nccwpck_require__(1527);
+async function run() {
+    await (0, saveImpl_1.default)(new stateProvider_1.StateProvider());
+}
+void run();
+exports["default"] = run;
+
+
+/***/ }),
+
+/***/ 6589:
+/***/ (function(__unused_webpack_module, exports, __nccwpck_require__) {
+
+"use strict";
+
 var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
     if (k2 === undefined) k2 = k;
     var desc = Object.getOwnPropertyDescriptor(m, k);
@@ -63003,15 +63038,17 @@ const core = __importStar(__nccwpck_require__(2186));
 const utils = __importStar(__nccwpck_require__(6850));
 const constants_1 = __nccwpck_require__(9042);
 process.on("uncaughtException", (error) => utils.logWarning(error.message));
-async function run() {
+async function saveImpl(stateProvider) {
+    let cacheId = -1;
     try {
         if (!utils.isCacheFeatureAvailable()) {
             return;
         }
-        const state = utils.getCacheState();
-        const primaryKey = core.getState(constants_1.State.CachePrimaryKey);
+        const state = stateProvider.getCacheState();
+        const primaryKey = stateProvider.getState(constants_1.State.CachePrimaryKey) ||
+            core.getInput(constants_1.Inputs.Key);
         if (!primaryKey) {
-            utils.logWarning(`Error retrieving key from state.`);
+            utils.logWarning(`Key is not specified.`);
             return;
         }
         if (utils.isExactKeyMatch(primaryKey, state)) {
@@ -63019,7 +63056,7 @@ async function run() {
             return;
         }
         const cachePaths = await utils.getBuildOutputPaths();
-        const cacheId = await cache.saveCache(cachePaths, primaryKey);
+        cacheId = await cache.saveCache(cachePaths, primaryKey);
         if (cacheId != -1) {
             core.info(`Cache saved with key: ${primaryKey}`);
         }
@@ -63027,8 +63064,58 @@ async function run() {
     catch (error) {
         utils.logWarning(error.message);
     }
+    return cacheId;
 }
-void run();
+exports["default"] = saveImpl;
+
+
+/***/ }),
+
+/***/ 1527:
+/***/ (function(__unused_webpack_module, exports, __nccwpck_require__) {
+
+"use strict";
+
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
+    __setModuleDefault(result, mod);
+    return result;
+};
+Object.defineProperty(exports, "__esModule", ({ value: true }));
+exports.StateProvider = void 0;
+const core = __importStar(__nccwpck_require__(2186));
+const constants_1 = __nccwpck_require__(9042);
+class StateProvider {
+    getCacheState() {
+        const cacheKey = this.getState(constants_1.State.CacheMatchedKey);
+        if (cacheKey) {
+            core.debug(`Cache state/key: ${cacheKey}`);
+            return cacheKey;
+        }
+        return undefined;
+    }
+    getState = core.getState;
+    setState = core.saveState;
+}
+exports.StateProvider = StateProvider;
 
 
 /***/ }),
@@ -63065,7 +63152,7 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", ({ value: true }));
-exports.getBuildOutputPaths = exports.getInputAsArray = exports.isCacheFeatureAvailable = exports.logWarning = exports.getCacheState = exports.setCacheHitOutput = exports.setCacheState = exports.setBuildMode = exports.isExactKeyMatch = exports.isGhes = void 0;
+exports.getBuildOutputPaths = exports.setBuildMode = exports.isCacheFeatureAvailable = exports.getInputAsArray = exports.isValidEvent = exports.logWarning = exports.isExactKeyMatch = exports.isGhes = void 0;
 const cache = __importStar(__nccwpck_require__(7799));
 const core = __importStar(__nccwpck_require__(2186));
 const path_1 = __importDefault(__nccwpck_require__(1017));
@@ -63083,31 +63170,23 @@ function isExactKeyMatch(key, cacheKey) {
         }) === 0);
 }
 exports.isExactKeyMatch = isExactKeyMatch;
-function setBuildMode(useCache) {
-    process.env.GATSBY_EXPERIMENTAL_PAGE_BUILD_ON_DATA_CHANGES = String(useCache);
-}
-exports.setBuildMode = setBuildMode;
-function setCacheState(state) {
-    core.saveState(constants_1.State.CacheMatchedKey, state);
-}
-exports.setCacheState = setCacheState;
-function setCacheHitOutput(isCacheHit) {
-    core.setOutput(constants_1.Outputs.CacheHit, isCacheHit.toString());
-}
-exports.setCacheHitOutput = setCacheHitOutput;
-function getCacheState() {
-    const cacheKey = core.getState(constants_1.State.CacheMatchedKey);
-    if (cacheKey) {
-        core.debug(`Cache state/key: ${cacheKey}`);
-        return cacheKey;
-    }
-    return undefined;
-}
-exports.getCacheState = getCacheState;
 function logWarning(message) {
-    core.info(`[warning]${message}`);
+    const warningPrefix = "[warning]";
+    core.info(`${warningPrefix}${message}`);
 }
 exports.logWarning = logWarning;
+function isValidEvent() {
+    return constants_1.RefKey in process.env && Boolean(process.env[constants_1.RefKey]);
+}
+exports.isValidEvent = isValidEvent;
+function getInputAsArray(name, options) {
+    return core
+        .getInput(name, options)
+        .split("\n")
+        .map((s) => s.replace(/^!\s+/, "!").trim())
+        .filter((x) => x !== "");
+}
+exports.getInputAsArray = getInputAsArray;
 function isCacheFeatureAvailable() {
     if (cache.isFeatureAvailable()) {
         return true;
@@ -63121,16 +63200,12 @@ Otherwise please upgrade to GHES version >= 3.5 and If you are also using Github
     return false;
 }
 exports.isCacheFeatureAvailable = isCacheFeatureAvailable;
-function getInputAsArray(name, options) {
-    return core
-        .getInput(name, options)
-        .split("\n")
-        .map((s) => s.replace(/^!\s+/, "!").trim())
-        .filter((x) => x !== "");
+function setBuildMode(useCache) {
+    process.env[constants_1.Gatsby.Env] = String(useCache);
 }
-exports.getInputAsArray = getInputAsArray;
+exports.setBuildMode = setBuildMode;
 async function getBuildOutputPaths() {
-    const targetPaths = [".cache", "public"];
+    const targetPaths = [constants_1.Gatsby.CacheDir, constants_1.Gatsby.PublicDir];
     const buildOutputPaths = [];
     for await (const target of targetPaths) {
         buildOutputPaths.push(path_1.default.join(process.cwd(), target));
diff --git a/src/constants.ts b/src/constants.ts
index d977635..70d33e7 100644
--- a/src/constants.ts
+++ b/src/constants.ts
@@ -12,3 +12,19 @@ export enum State {
   CachePrimaryKey = "CACHE_KEY",
   CacheMatchedKey = "CACHE_RESULT",
 }
+
+export enum Events {
+  Key = "GITHUB_EVENT_NAME",
+}
+
+export enum Platform {
+  RunnerOs = "RUNNER_OS",
+}
+
+export enum Gatsby {
+  CacheDir = ".cache",
+  PublicDir = "public",
+  Env = "GATSBY_EXPERIMENTAL_PAGE_BUILD_ON_DATA_CHANGES",
+}
+
+export const RefKey = "GITHUB_REF";
diff --git a/src/restore.ts b/src/restore.ts
index 30a90b9..7b63f9d 100644
--- a/src/restore.ts
+++ b/src/restore.ts
@@ -1,58 +1,10 @@
-import * as cache from "@actions/cache";
-import * as core from "@actions/core";
-import * as utils from "./utils/actionUtils";
-import { Inputs, State } from "./constants";
-
-type Platform = "Linux" | "Windows" | "macOS";
+import restoreImpl from "./restoreImpl";
+import { StateProvider } from "./stateProvider";
 
 async function run(): Promise<void> {
-  try {
-    if (!utils.isCacheFeatureAvailable()) {
-      utils.setCacheHitOutput(false);
-      return;
-    }
-
-    const useCache: boolean = core.getBooleanInput(Inputs.UseCache);
-    utils.setBuildMode(useCache);
-
-    const cachePaths: string[] = await utils.getBuildOutputPaths();
-    const restoreKeys: string[] = utils.getInputAsArray(Inputs.RestoreKeys);
-
-    let primaryKey: string = core.getInput(Inputs.Key);
-    if (!primaryKey) {
-      const platform: Platform = process.env.RUNNER_OS as Platform;
-      primaryKey = `${platform}-gatsby-build-`;
-    }
-
-    core.debug(`primary key is ${primaryKey}`);
-    core.saveState(State.CachePrimaryKey, primaryKey);
-
-    const cacheKey: string | undefined = await cache.restoreCache(
-      cachePaths,
-      primaryKey,
-      restoreKeys
-    );
-
-    if (!cacheKey) {
-      core.info(
-        `Cache not found for keys: ${[primaryKey, ...restoreKeys].join(", ")}`
-      );
-
-      return;
-    }
-
-    utils.setCacheState(cacheKey);
-
-    const isExactKeyMatch: boolean = utils.isExactKeyMatch(
-      primaryKey,
-      cacheKey
-    );
-
-    utils.setCacheHitOutput(isExactKeyMatch);
-    core.info(`Cache restored from key: ${cacheKey}`);
-  } catch (error: unknown) {
-    core.setFailed((error as Error).message);
-  }
+  await restoreImpl(new StateProvider());
 }
 
 void run();
+
+export default run;
diff --git a/src/restoreImpl.ts b/src/restoreImpl.ts
new file mode 100644
index 0000000..77b49c2
--- /dev/null
+++ b/src/restoreImpl.ts
@@ -0,0 +1,72 @@
+import * as cache from "@actions/cache";
+import * as core from "@actions/core";
+import * as utils from "./utils/actionUtils";
+import { Events, Inputs, Outputs, Platform, State } from "./constants";
+import { StateProvider } from "./stateProvider";
+
+type Runner = "Linux" | "Windows" | "macOS";
+
+async function restoreImpl(
+  stateProvider: StateProvider
+): Promise<string | undefined> {
+  try {
+    if (!utils.isCacheFeatureAvailable()) {
+      core.setOutput(Outputs.CacheHit, "false");
+      return;
+    }
+
+    if (!utils.isValidEvent()) {
+      const eventName: string = process.env[Events.Key] || "";
+      utils.logWarning(
+        `Event Validation Error: The event type ${eventName} is not supported because it's not tied to a branch or tag ref.`
+      );
+      return;
+    }
+
+    const useCache: boolean = core.getBooleanInput(Inputs.UseCache);
+    utils.setBuildMode(useCache);
+
+    const cachePaths: string[] = await utils.getBuildOutputPaths();
+    const restoreKeys: string[] = utils.getInputAsArray(Inputs.RestoreKeys);
+
+    let primaryKey: string = core.getInput(Inputs.Key);
+    if (!primaryKey) {
+      const platform: Runner = process.env[Platform.RunnerOs] as Runner;
+      primaryKey = `${platform}-gatsby-build-`;
+    }
+
+    core.debug(`primary key is ${primaryKey}`);
+    stateProvider.setState(State.CachePrimaryKey, primaryKey);
+
+    const cacheKey: string | undefined = await cache.restoreCache(
+      cachePaths,
+      primaryKey,
+      restoreKeys
+    );
+
+    if (!cacheKey) {
+      core.info(
+        `Cache not found for keys: ${[primaryKey, ...restoreKeys].join(", ")}`
+      );
+
+      return;
+    }
+
+    stateProvider.setState(State.CacheMatchedKey, cacheKey);
+
+    const isExactKeyMatch: boolean = utils.isExactKeyMatch(
+      primaryKey,
+      cacheKey
+    );
+
+    core.setOutput(Outputs.CacheHit, isExactKeyMatch.toString());
+    core.info(`Cache restored from key: ${cacheKey}`);
+
+    return cacheKey;
+  } catch (error: unknown) {
+    core.setFailed((error as Error).message);
+    return;
+  }
+}
+
+export default restoreImpl;
diff --git a/src/save.ts b/src/save.ts
index 88978d5..8c96a6c 100644
--- a/src/save.ts
+++ b/src/save.ts
@@ -1,42 +1,10 @@
-import * as cache from "@actions/cache";
-import * as core from "@actions/core";
-import * as utils from "./utils/actionUtils";
-import { State } from "./constants";
-
-process.on("uncaughtException", (error: unknown) =>
-  utils.logWarning((error as Error).message)
-);
+import saveImpl from "./saveImpl";
+import { StateProvider } from "./stateProvider";
 
 async function run(): Promise<void> {
-  try {
-    if (!utils.isCacheFeatureAvailable()) {
-      return;
-    }
-
-    const state: string | undefined = utils.getCacheState();
-
-    const primaryKey: string = core.getState(State.CachePrimaryKey);
-    if (!primaryKey) {
-      utils.logWarning(`Error retrieving key from state.`);
-      return;
-    }
-
-    if (utils.isExactKeyMatch(primaryKey, state)) {
-      core.info(
-        `Cache hit occurred on the primary key ${primaryKey}, not saving cache.`
-      );
-      return;
-    }
-
-    const cachePaths: string[] = await utils.getBuildOutputPaths();
-    const cacheId: number = await cache.saveCache(cachePaths, primaryKey);
-
-    if (cacheId != -1) {
-      core.info(`Cache saved with key: ${primaryKey}`);
-    }
-  } catch (error: unknown) {
-    utils.logWarning((error as Error).message);
-  }
+  await saveImpl(new StateProvider());
 }
 
 void run();
+
+export default run;
diff --git a/src/saveImpl.ts b/src/saveImpl.ts
new file mode 100644
index 0000000..ef0ec21
--- /dev/null
+++ b/src/saveImpl.ts
@@ -0,0 +1,50 @@
+import * as cache from "@actions/cache";
+import * as core from "@actions/core";
+import * as utils from "./utils/actionUtils";
+import { Inputs, State } from "./constants";
+import { BaseStateProvider } from "./stateProvider";
+
+process.on("uncaughtException", (error: unknown) =>
+  utils.logWarning((error as Error).message)
+);
+
+async function saveImpl(
+  stateProvider: BaseStateProvider
+): Promise<number | void> {
+  let cacheId = -1;
+  try {
+    if (!utils.isCacheFeatureAvailable()) {
+      return;
+    }
+
+    const state: string | undefined = stateProvider.getCacheState();
+
+    const primaryKey: string =
+      stateProvider.getState(State.CachePrimaryKey) ||
+      core.getInput(Inputs.Key);
+
+    if (!primaryKey) {
+      utils.logWarning(`Key is not specified.`);
+      return;
+    }
+
+    if (utils.isExactKeyMatch(primaryKey, state)) {
+      core.info(
+        `Cache hit occurred on the primary key ${primaryKey}, not saving cache.`
+      );
+      return;
+    }
+
+    const cachePaths: string[] = await utils.getBuildOutputPaths();
+    cacheId = await cache.saveCache(cachePaths, primaryKey);
+
+    if (cacheId != -1) {
+      core.info(`Cache saved with key: ${primaryKey}`);
+    }
+  } catch (error: unknown) {
+    utils.logWarning((error as Error).message);
+  }
+  return cacheId;
+}
+
+export default saveImpl;
diff --git a/src/stateProvider.ts b/src/stateProvider.ts
new file mode 100644
index 0000000..d804f5e
--- /dev/null
+++ b/src/stateProvider.ts
@@ -0,0 +1,23 @@
+import * as core from "@actions/core";
+import { State } from "./constants";
+
+export interface BaseStateProvider {
+  getCacheState(): string | undefined;
+  getState(key: string): string;
+  setState(key: string, value: string): void;
+}
+
+export class StateProvider implements BaseStateProvider {
+  getCacheState(): string | undefined {
+    const cacheKey: string = this.getState(State.CacheMatchedKey);
+    if (cacheKey) {
+      core.debug(`Cache state/key: ${cacheKey}`);
+      return cacheKey;
+    }
+
+    return undefined;
+  }
+
+  getState = core.getState;
+  setState = core.saveState;
+}
diff --git a/src/utils/actionUtils.ts b/src/utils/actionUtils.ts
index acc8c13..474efb0 100644
--- a/src/utils/actionUtils.ts
+++ b/src/utils/actionUtils.ts
@@ -1,7 +1,7 @@
 import * as cache from "@actions/cache";
 import * as core from "@actions/core";
 import path from "path";
-import { Outputs, State } from "../constants";
+import { Gatsby, RefKey } from "../constants";
 
 export function isGhes(): boolean {
   const url: string = process.env.GITHUB_SERVER_URL ?? "https://github.com";
@@ -18,30 +18,24 @@ export function isExactKeyMatch(key: string, cacheKey?: string): boolean {
   );
 }
 
-export function setBuildMode(useCache: boolean): void {
-  process.env.GATSBY_EXPERIMENTAL_PAGE_BUILD_ON_DATA_CHANGES = String(useCache);
-}
-
-export function setCacheState(state: string): void {
-  core.saveState(State.CacheMatchedKey, state);
-}
-
-export function setCacheHitOutput(isCacheHit: boolean): void {
-  core.setOutput(Outputs.CacheHit, isCacheHit.toString());
+export function logWarning(message: string): void {
+  const warningPrefix = "[warning]";
+  core.info(`${warningPrefix}${message}`);
 }
 
-export function getCacheState(): string | undefined {
-  const cacheKey: string = core.getState(State.CacheMatchedKey);
-  if (cacheKey) {
-    core.debug(`Cache state/key: ${cacheKey}`);
-    return cacheKey;
-  }
-
-  return undefined;
+export function isValidEvent(): boolean {
+  return RefKey in process.env && Boolean(process.env[RefKey]);
 }
 
-export function logWarning(message: string): void {
-  core.info(`[warning]${message}`);
+export function getInputAsArray(
+  name: string,
+  options?: core.InputOptions
+): string[] {
+  return core
+    .getInput(name, options)
+    .split("\n")
+    .map((s) => s.replace(/^!\s+/, "!").trim())
+    .filter((x) => x !== "");
 }
 
 export function isCacheFeatureAvailable(): boolean {
@@ -63,19 +57,12 @@ Otherwise please upgrade to GHES version >= 3.5 and If you are also using Github
   return false;
 }
 
-export function getInputAsArray(
-  name: string,
-  options?: core.InputOptions
-): string[] {
-  return core
-    .getInput(name, options)
-    .split("\n")
-    .map((s) => s.replace(/^!\s+/, "!").trim())
-    .filter((x) => x !== "");
+export function setBuildMode(useCache: boolean): void {
+  process.env[Gatsby.Env] = String(useCache);
 }
 
 export async function getBuildOutputPaths(): Promise<string[]> {
-  const targetPaths: string[] = [".cache", "public"];
+  const targetPaths: string[] = [Gatsby.CacheDir, Gatsby.PublicDir];
   const buildOutputPaths: string[] = [];
 
   for await (const target of targetPaths) {