fix(yoga): use the onContextBuilding hook

Use the onContextBuilding to load the data into an existing federated
token instead of the onRequest variant which created one if it didnt
exist.

This moves the initial creation of the federated token to the
implementation side and thus more straightforward

.changeset/fluffy-rivers-explode.md

@@ -0,0 +1,5 @@
+---
+"@labdigital/federated-token-yoga": minor
+---
+
+Require the FederatedToken to already exist on the context

.eslintrc.cjs

@@ -12,7 +12,6 @@ module.exports = {
 			},
 		],
 		"unused-imports/no-unused-imports": "error",
-		"arrow-body-style": ["error", "as-needed"],
 		"no-console": [
 			"error",
 			{

packages/yoga/src/index.test.ts

@@ -0,0 +1,147 @@
+import { createSchema, createYoga } from "graphql-yoga";
+import { describe, it, expect } from "vitest";
+import { useFederatedToken } from ".";
+import { FederatedToken } from "@labdigital/federated-token";
+
+type FederatedTokenContext = {
+	federatedToken: FederatedToken;
+};
+
+export function assertSingleValue<TValue extends object>(
+	value: TValue | AsyncIterable<TValue>,
+): asserts value is TValue {
+	if (Symbol.asyncIterator in value) {
+		throw new Error("Expected single value");
+	}
+}
+
+describe("useFederatdToken()", () => {
+	const schemaFactory = async () => {
+		return createSchema<FederatedTokenContext>({
+			typeDefs: /* GraphQL */ `
+				type Query {
+					tokenData: String
+				}
+				type Mutation {
+					createToken(service: String!): Boolean
+				}
+			`,
+			resolvers: {
+				Query: {
+					tokenData: (parent, args, context) => {
+						return JSON.stringify({
+							tokens: context.federatedToken.tokens,
+							values: context.federatedToken.values,
+							refreshTokens: context.federatedToken.refreshTokens,
+						});
+					},
+				},
+				Mutation: {
+					createToken: (parent, args, context) => {
+						context.federatedToken.setAccessToken(args.service, {
+							token: "test",
+							exp: 123,
+							sub: "test",
+						});
+						context.federatedToken.setRefreshToken(
+							args.service,
+							"refresh-token",
+						);
+						return true;
+					},
+				},
+			},
+		});
+	};
+
+	const yoga = createYoga<FederatedTokenContext>({
+		schema: schemaFactory,
+		plugins: [useFederatedToken()],
+		context: () => {
+			return {
+				federatedToken: new FederatedToken(),
+			};
+		},
+	});
+
+	it("should handle missing tokens", async () => {
+		const response = await yoga.fetch("http://localhost:3000/graphql", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+			},
+			body: JSON.stringify({
+				query: "{ tokenData }",
+			}),
+		});
+		const content = await response.json();
+		assertSingleValue(content);
+	});
+
+	it("should load access and refresh tokens", async () => {
+		const token = new FederatedToken();
+		token.setValue("my-service", { foo: "bar" });
+		token.setAccessToken("my-service", { token: "bar", exp: 123, sub: "test" });
+		token.setRefreshToken("my-service", "refresh-token");
+
+		const response = await yoga.fetch("http://localhost:3000/graphql", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+				"X-Access-Token": token.serializeAccessToken() ?? "",
+				"X-Refresh-Token": token.dumpRefreshToken() ?? "",
+			},
+			body: JSON.stringify({
+				query: "{ tokenData }",
+			}),
+		});
+		const content = await response.json();
+
+		assertSingleValue(content);
+
+		const result = JSON.parse(content.data.tokenData);
+		expect(result).toEqual({
+			tokens: {
+				"my-service": {
+					token: "bar",
+					exp: 123,
+					sub: "test",
+				},
+			},
+			values: {
+				"my-service": {
+					foo: "bar",
+				},
+			},
+			refreshTokens: {
+				"my-service": "refresh-token",
+			},
+		});
+
+		expect(response.headers.has("X-Access-Token")).toBe(false);
+		expect(response.headers.has("X-Refresh-Token")).toBe(false);
+	});
+
+	it("should create new token", async () => {
+		const response = await yoga.fetch("http://localhost:3000/graphql", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+			},
+			body: JSON.stringify({
+				query: `mutation { createToken(service: "foobar") }`,
+			}),
+		});
+		const content = await response.json();
+		assertSingleValue(content);
+
+		expect(response.headers.has("X-Access-Token")).toBe(true);
+		expect(response.headers.has("X-Refresh-Token")).toBe(true);
+
+		const accessToken = response.headers.get("X-Access-Token");
+		const refreshToken = response.headers.get("X-Refresh-Token");
+
+		expect(accessToken).toBeDefined();
+		expect(refreshToken).toBeDefined();
+	});
+});

packages/yoga/src/index.ts

@@ -6,26 +6,26 @@ type FederatedTokenContext = {
 };
 
 export const useFederatedToken = <T extends FederatedTokenContext>(): Plugin<
-	object,
+	T,
 	T
 > => ({
-	onRequest: ({ request, serverContext }) => {
-		// Initialize FederatedToken and add it to the context
-		const federatedToken = serverContext.federatedToken ?? new FederatedToken();
-
-		// Retrieve tokens from headers using the serverContext
-		const accessToken = request.headers.get("x-access-token") as string;
-		const refreshToken = request.headers.get("x-refresh-token") as string;
+	onContextBuilding: ({ context, extendContext }) => {
+		if (!context.federatedToken) {
+			throw new Error("Federated token not found in context");
+		}
+		if (!context.request) {
+			throw new Error("Request not found in context");
+		}
 
+		const { request, federatedToken } = context;
+		const accessToken = request.headers.get("x-access-token");
+		const refreshToken = request.headers.get("x-refresh-token");
 		if (accessToken) {
 			federatedToken.deserializeAccessToken(accessToken);
 		}
-
 		if (refreshToken) {
 			federatedToken.loadRefreshToken(refreshToken);
 		}
-
-		serverContext.federatedToken = federatedToken;
 	},
 
 	onResponse: async ({ response, serverContext }) => {