Merge pull request #28 from upstash/DX-1393-auth-in-failure-function

Custom auth in failure function

examples/ci/app/ci/ci.test.ts

@@ -10,7 +10,7 @@ describe("workflow integration tests", () => {
         await initiateTest(testConfig.route, testConfig.waitForSeconds)
       },
       {
-        timeout: (testConfig.waitForSeconds + 10) * 1000
+        timeout: (testConfig.waitForSeconds + 15) * 1000
       }
     )
   });

examples/ci/app/ci/constants.ts

@@ -31,14 +31,19 @@ export const TEST_ROUTES: Pick<TestConfig, "route" | "waitForSeconds">[] = [
   },
   {
     // checks auth
-    route: "auth",
+    route: "auth/success",
     waitForSeconds: 1
   },
   {
     // checks auth failing
-    route: "auth-fail",
+    route: "auth/fail",
     waitForSeconds: 0
   },
+  {
+    // checks custom auth
+    route: "auth/custom/workflow",
+    waitForSeconds: 5
+  },
   {
     // checks context.call (sucess and fail case)
     route: "call/workflow",

examples/ci/app/ci/upstash/redis.test.ts

@@ -80,4 +80,21 @@ describe("redis", () => {
       ).not.toThrow()
     })
   })
+
+  test("should fail if marked as failed", async () => {
+
+    const route = "fail-route"
+    const randomId = `random-id-${nanoid()}`
+    const result = `random-result-${nanoid()}`
+
+    // increment, save and check
+    await redis.increment(route, randomId)
+    await redis.saveResultsWithoutContext(route, randomId, result)
+    await redis.checkRedisForResults(route, randomId, 1, result)
+
+    // mark as failed and check
+    await redis.failWithoutContext(route, randomId)
+    expect(redis.checkRedisForResults(route, randomId, 1, result)).rejects.toThrow(redis.FAILED_TEXT)
+
+  })
 })

examples/ci/app/ci/upstash/redis.ts

@@ -8,7 +8,7 @@ const redis = Redis.fromEnv();
 const EXPIRE_IN_SECS = 60
 
 const getRedisKey = (
-  kind: "increment" | "result",
+  kind: "increment" | "result" | "fail",
   route: string,
   randomTestId: string
 ): string => {
@@ -46,11 +46,7 @@ export const saveResultsWithoutContext = async (
 
   // save result
   const key = getRedisKey("result", route, randomTestId)
-
-  const pipe = redis.pipeline()
-  pipe.set<RedisResult>(key, { callCount, result, randomTestId })
-  pipe.expire(key, EXPIRE_IN_SECS)
-  await pipe.exec()
+  await redis.set<RedisResult>(key, { callCount, result, randomTestId }, { ex: EXPIRE_IN_SECS })
 }
 
 /**
@@ -80,6 +76,38 @@ export const saveResult = async (
   )
 }
 
+export const failWithoutContext = async (
+  route: string,
+  randomTestId: string
+) => {
+  const key = getRedisKey("fail", route, randomTestId)
+  await redis.set<boolean>(key, true, { ex: EXPIRE_IN_SECS })
+}
+
+/**
+ * marks the workflow as failed
+ * 
+ * @param context 
+ * @returns 
+ */
+export const fail = async (
+  context: WorkflowContext<unknown>,
+) => {
+  const randomTestId = context.headers.get(CI_RANDOM_ID_HEADER)
+  const route = context.headers.get(CI_ROUTE_HEADER)
+
+  if (randomTestId === null) {
+    throw new Error("randomTestId can't be null.")
+  }
+  if (route === null) {
+    throw new Error("route can't be null.")
+  }
+
+  await failWithoutContext(route, randomTestId)
+}
+
+export const FAILED_TEXT = "Test has failed because it was marked as failed with `fail` method."
+
 export const checkRedisForResults = async (
   route: string,
   randomTestId: string,
@@ -101,6 +129,12 @@ export const checkRedisForResults = async (
     throw new Error(`result not found for route ${route} with randomTestId ${randomTestId}`)
   }
 
+  const failKey = getRedisKey("fail", route, randomTestId)
+  const failed = await redis.get<boolean>(failKey)
+  if (failed) {
+    throw new Error(FAILED_TEXT)
+  }
+
   const { callCount, randomTestId: resultRandomTestId, result } = testResult 
 
   expect(resultRandomTestId, randomTestId)

examples/ci/app/test-routes/auth/README.md

@@ -0,0 +1,4 @@
+this directory has three tests
+- success: checking auth correctly
+- fail: auth failing
+- custom: define an workflow endpoint secured with custom auth (instead of receiver) and try to call it as if failure callback

examples/ci/app/test-routes/auth/custom/target/route.ts

@@ -0,0 +1,16 @@
+import { WorkflowContext } from "@upstash/workflow";
+import { serve } from "@upstash/workflow/nextjs";
+import { fail } from "app/ci/upstash/redis";
+import { nanoid } from "app/ci/utils";
+
+
+export const { POST } = serve(async (context) => {
+  if (context.headers.get("authorization") !== nanoid()) {
+    return;
+  };
+}, {
+  receiver: undefined,
+  async failureFunction({ context }) {
+    await fail(context as WorkflowContext)
+  },
+})

examples/ci/app/test-routes/auth/custom/workflow/route.ts

@@ -0,0 +1,95 @@
+import { serve } from "@upstash/workflow/nextjs";
+import { BASE_URL, CI_RANDOM_ID_HEADER, CI_ROUTE_HEADER, TEST_ROUTE_PREFIX } from "app/ci/constants";
+import { testServe, expect } from "app/ci/utils";
+import { FailureFunctionPayload, WorkflowContext } from "@upstash/workflow";
+import { saveResult } from "app/ci/upstash/redis";
+
+const header = `test-header-foo`
+const headerValue = `header-bar`
+const authentication = `Bearer test-auth-super-secret`
+const payload = "my-payload"
+
+const thirdPartyEndpoint = `${TEST_ROUTE_PREFIX}/auth/custom/target`
+
+const makeCall = async (
+  context: WorkflowContext,
+  stepName: string,
+  method: "GET" | "POST",
+  expectedStatus: number,
+  expectedBody: unknown
+) => {
+  const randomId = context.headers.get(CI_RANDOM_ID_HEADER)
+  const route = context.headers.get(CI_ROUTE_HEADER)
+
+  if (!randomId || !route) {
+    throw new Error("randomId or route not found")
+  }
+
+  const { status, body } = await context.call<FailureFunctionPayload>(stepName, {
+    url: thirdPartyEndpoint,
+    body: 
+    {
+      status: 200,
+      header: "",
+      body: "",
+      url: "",
+      sourceHeader: {
+        [CI_ROUTE_HEADER]: [route],
+        [CI_RANDOM_ID_HEADER]: [randomId]
+      },
+      sourceBody: "",
+      workflowRunId: "",
+      sourceMessageId: "",
+    },
+    method,
+    headers: {
+      [ CI_RANDOM_ID_HEADER ]: randomId,
+      [ CI_ROUTE_HEADER ]: route,
+      "Upstash-Workflow-Is-Failure": "true"
+    }
+  })
+
+  expect(status, expectedStatus)
+
+  expect(typeof body, typeof expectedBody)
+  expect(JSON.stringify(body), JSON.stringify(expectedBody))
+}
+
+export const { POST, GET } = testServe(
+  serve<string>(
+    async (context) => {
+
+      expect(context.headers.get(header)!, headerValue)
+
+      await makeCall(
+        context,
+        "regular call should fail",
+        "POST",
+        500,
+        {
+          error: "WorkflowError",
+          message: "Not authorized to run the failure function."
+        }
+      )
+
+      const input = context.requestPayload;
+      expect(input, payload);
+
+      await saveResult(
+        context,
+        "not authorized for failure"
+      )
+    }, {
+      baseUrl: BASE_URL,
+      retries: 0,
+    }
+  ), {
+    expectedCallCount: 4,
+    expectedResult: "not authorized for failure",
+    payload,
+    headers: {
+      [ header ]: headerValue,
+      "authentication": authentication
+    }
+  }
+) 

examples/ci/app/test-routes/auth-fail/route.ts → examples/ci/app/test-routes/auth/fail/route.ts

@@ -1,7 +1,7 @@
 import { serve } from "@upstash/workflow/nextjs";
 import { BASE_URL } from "app/ci/constants";
 import { testServe, expect } from "app/ci/utils";
-import { saveResult } from "app/ci/upstash/redis"
+import { fail, saveResult } from "app/ci/upstash/redis"
 
 const header = `test-header-foo`
 const headerValue = `header-bar`
@@ -28,10 +28,10 @@ export const { POST, GET } = testServe(
         return;
       }
 
-      throw new Error("shouldn't come here.")
+      await fail(context)
     }, {
       baseUrl: BASE_URL,
-      retries: 0
+      retries: 1 // check with retries 1 to see if endpoint will retry
     }
   ), {
     expectedCallCount: 1,

examples/ci/app/test-routes/call/third-party/route.ts

@@ -30,5 +30,15 @@ export const PATCH = async () => {
       headers: {
         [ FAILING_HEADER ]: FAILING_HEADER_VALUE
       }
-    })
-}
+    }
+  )
+}
+
+export const PUT = async () => {
+  return new Response(
+    undefined,
+    {
+      status: 300,
+    }
+  )
+}

examples/ci/app/test-routes/call/workflow-with-failureFunction/route.ts

@@ -1,8 +1,9 @@
 import { serve } from "@upstash/workflow/nextjs";
 import { BASE_URL, TEST_ROUTE_PREFIX } from "app/ci/constants";
 import { testServe, expect } from "app/ci/utils";
-import { saveResult } from "app/ci/upstash/redis"
+import { fail, saveResult } from "app/ci/upstash/redis"
 import { FAILING_HEADER, FAILING_HEADER_VALUE } from "../constants";
+import { WorkflowContext } from "@upstash/workflow";
 
 const testHeader = `test-header-foo`
 const headerValue = `header-foo`
@@ -31,8 +32,8 @@ export const { POST, GET } = testServe(
     }, {
       baseUrl: BASE_URL,
       retries: 0,
-      failureFunction() {
-        console.log("SHOULDNT RUN");
+      async failureFunction({ context }) {
+        await fail(context as WorkflowContext)
       },
     }
   ), {

examples/ci/app/test-routes/call/workflow/route.ts

@@ -45,16 +45,16 @@ export const { POST, GET } = testServe(
 
       await context.sleep("sleep 1", 2);
 
-      const { body: getResult, header: getHeaders, status: getStatus } = await context.call("get call", {
+      const { body: getResult, header: getHeaders, status: getStatus } = await context.call<string>("get call", {
         url: thirdPartyEndpoint,
         headers: getHeader,
       });
 
       expect(getStatus, 200)
       expect(getHeaders[GET_HEADER][0], GET_HEADER_VALUE)
-      expect(getResult as string, "called GET 'third-party-result' 'get-header-value-x'");
+      expect(getResult, "called GET 'third-party-result' 'get-header-value-x'");
 
-      const { body: patchResult, status, header } = await context.call("get call", {
+      const { body: patchResult, status, header } = await context.call("patch call", {
         url: thirdPartyEndpoint,
         headers: getHeader,
         method: "PATCH",
@@ -65,16 +65,26 @@ export const { POST, GET } = testServe(
       expect(patchResult as string, "failing request");
       expect(header[FAILING_HEADER][0], FAILING_HEADER_VALUE)
 
+      // put will return with an empty body. should return "" as body in that case.
+      const { body: putBody, status: putStatus } = await context.call<string>("put call", {
+        url: thirdPartyEndpoint,
+        method: "PUT",
+        retries: 0
+      })
+
+      expect(putStatus, 300)
+      expect(putBody, "");
+
       await saveResult(
         context,
-        getResult as string
+        getResult
       )
     }, {
       baseUrl: BASE_URL,
       retries: 0
     }
   ), {
-    expectedCallCount: 10,
+    expectedCallCount: 12,
     expectedResult: "called GET 'third-party-result' 'get-header-value-x'",
     payload,
     headers: {

examples/cloudflare-workers-hono/ci.test.ts

@@ -96,7 +96,7 @@ const testEndpoint = ({
     expect(result?.secret).toBe(secret)
     expect(result?.result).toBe(expectedResult)
   }, {
-    timeout: 8000
+    timeout: 15000
   })
 }
 

examples/nextjs-pages/ci.test.ts

@@ -107,7 +107,7 @@ const testEndpoint = ({
     expect(result?.secret).toBe(secret)
     expect(result?.result).toBe(expectedResult)
   }, {
-    timeout: 9000
+    timeout: 15000
   })
 }
 

src/serve/index.ts

@@ -97,6 +97,7 @@ export const serve = <
       requestPayload,
       qstashClient,
       initialPayloadParser,
+      routeFunction,
       failureFunction
     );
     if (failureCheck.isErr()) {