From 09c03d46b0c52a40d58ff628b49abd60389097df Mon Sep 17 00:00:00 2001
From: Mwelwa Nkuta <mwelwankuta@gmail.com>
Date: Wed, 1 Feb 2023 00:01:29 +0200
Subject: [PATCH] Fix for #18

---
 pages/api/generate.ts | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/pages/api/generate.ts b/pages/api/generate.ts
index 25d97d2..5356b1d 100644
--- a/pages/api/generate.ts
+++ b/pages/api/generate.ts
@@ -10,6 +10,10 @@ interface ExtendedNextApiRequest extends NextApiRequest {
   };
 }
 
+interface EndPointSuccess {
+  pass: string | number | null;
+  fail: string | number | null;
+}
 // Create a new ratelimiter, that allows 3 requests per 60 seconds
 const ratelimit = redis
   ? new Ratelimit({
@@ -18,6 +22,14 @@ const ratelimit = redis
     })
   : undefined;
 
+const endpointResponse = (restoredImage: string | null, endpointSuccess: EndPointSuccess) => {
+  const {pass, fail} = endpointSuccess;
+  if (restoredImage) {
+    return pass;
+  }
+  return fail;
+}
+
 export default async function handler(
   req: ExtendedNextApiRequest,
   res: NextApiResponse<Data>
@@ -79,7 +91,9 @@ export default async function handler(
       await new Promise((resolve) => setTimeout(resolve, 1000));
     }
   }
+  const status = endpointResponse(restoredImage, {pass: 200, fail: 500})
+  const json = endpointResponse(restoredImage, {pass: restoredImage, fail: "Failed to restore image"})
   res
-    .status(200)
-    .json(restoredImage ? restoredImage : "Failed to restore image");
+    .status(status)
+    .json(json);
 }