From fedd9e04d8dd82b8eb4d1f8dac57ff65c0a96a6f Mon Sep 17 00:00:00 2001
From: Joachim Van Herwegen <joachimvh@gmail.com>
Date: Wed, 4 Oct 2023 16:17:53 +0200
Subject: [PATCH] feat: Allow ConditionalHandler to set the expected value

---
 src/util/handlers/ConditionalHandler.ts       | 21 ++++++++++++++++---
 .../util/handlers/ConditionalHandler.test.ts  | 18 ++++++++++++++++
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/src/util/handlers/ConditionalHandler.ts b/src/util/handlers/ConditionalHandler.ts
index 5fd4ec6246..ddf4229752 100644
--- a/src/util/handlers/ConditionalHandler.ts
+++ b/src/util/handlers/ConditionalHandler.ts
@@ -8,22 +8,27 @@ import { AsyncHandler } from './AsyncHandler';
  * After that all input will be rejected.
  * Once the value has been matched this behaviour will be cached,
  * so changing the value again afterwards will not enable this handler again.
+ *
+ * If `handleStorage` is set to `true`,
+ * this handler will set the value itself in the given storage after the source handler successfully resolved.
  */
 export class ConditionalHandler<TIn, TOut> extends AsyncHandler<TIn, TOut> {
   private readonly source: AsyncHandler<TIn, TOut>;
   private readonly storage: KeyValueStorage<string, unknown>;
   private readonly storageKey: string;
   private readonly storageValue: unknown;
+  private readonly handleStorage: boolean;
 
   private finished: boolean;
 
   public constructor(source: AsyncHandler<TIn, TOut>, storage: KeyValueStorage<string, unknown>, storageKey: string,
-    storageValue: unknown) {
+    storageValue: unknown, handleStorage = false) {
     super();
     this.source = source;
     this.storage = storage;
     this.storageKey = storageKey;
     this.storageValue = storageValue;
+    this.handleStorage = handleStorage;
 
     this.finished = false;
   }
@@ -35,11 +40,21 @@ export class ConditionalHandler<TIn, TOut> extends AsyncHandler<TIn, TOut> {
 
   public async handleSafe(input: TIn): Promise<TOut> {
     await this.checkCondition();
-    return this.source.handleSafe(input);
+    const result = await this.source.handleSafe(input);
+    if (this.handleStorage) {
+      await this.storage.set(this.storageKey, this.storageValue);
+      this.finished = true;
+    }
+    return result;
   }
 
   public async handle(input: TIn): Promise<TOut> {
-    return this.source.handle(input);
+    const result = await this.source.handle(input);
+    if (this.handleStorage) {
+      await this.storage.set(this.storageKey, this.storageValue);
+      this.finished = true;
+    }
+    return result;
   }
 
   /**
diff --git a/test/unit/util/handlers/ConditionalHandler.test.ts b/test/unit/util/handlers/ConditionalHandler.test.ts
index 5e6ef17dd0..8f76ab9d41 100644
--- a/test/unit/util/handlers/ConditionalHandler.test.ts
+++ b/test/unit/util/handlers/ConditionalHandler.test.ts
@@ -45,6 +45,15 @@ describe('A ConditionalHandler', (): void => {
     await expect(handler.handleSafe(input)).resolves.toBe('handledSafely');
     expect(source.handleSafe).toHaveBeenCalledTimes(1);
     expect(source.handleSafe).toHaveBeenLastCalledWith(input);
+    expect(storage.get(storageKey)).toBeUndefined();
+  });
+
+  it('can store the value itself if requested after a handleSafe call.', async(): Promise<void> => {
+    handler = new ConditionalHandler(source, storage, storageKey, storageValue, true);
+    await expect(handler.handleSafe(input)).resolves.toBe('handledSafely');
+    expect(source.handleSafe).toHaveBeenCalledTimes(1);
+    expect(source.handleSafe).toHaveBeenLastCalledWith(input);
+    expect(storage.get(storageKey)).toBe(storageValue);
   });
 
   it('rejects all handleSafe requests once the storage value matches.', async(): Promise<void> => {
@@ -57,6 +66,7 @@ describe('A ConditionalHandler', (): void => {
     await expect(handler.handle(input)).resolves.toBe('handled');
     expect(source.handle).toHaveBeenCalledTimes(1);
     expect(source.handle).toHaveBeenLastCalledWith(input);
+    expect(storage.get(storageKey)).toBeUndefined();
   });
 
   it('does not reject handle requests once the storage value matches.', async(): Promise<void> => {
@@ -65,4 +75,12 @@ describe('A ConditionalHandler', (): void => {
     expect(source.handle).toHaveBeenCalledTimes(1);
     expect(source.handle).toHaveBeenLastCalledWith(input);
   });
+
+  it('can store the value itself if requested after a handle call.', async(): Promise<void> => {
+    handler = new ConditionalHandler(source, storage, storageKey, storageValue, true);
+    await expect(handler.handle(input)).resolves.toBe('handled');
+    expect(source.handle).toHaveBeenCalledTimes(1);
+    expect(source.handle).toHaveBeenLastCalledWith(input);
+    expect(storage.get(storageKey)).toBe(storageValue);
+  });
 });