From 79a3328d752a3efb7558876a64b2b0e367d89714 Mon Sep 17 00:00:00 2001
From: Alex Gherghisan <alexg@aztecprotocol.com>
Date: Thu, 18 Jan 2024 08:48:33 +0000
Subject: [PATCH] test: add tests

---
 .../kv-store/src/lmdb/counter.test.ts         | 115 ++++++++++++++++++
 yarn-project/kv-store/src/lmdb/counter.ts     |   9 +-
 yarn-project/kv-store/src/lmdb/store.ts       |   4 +-
 3 files changed, 125 insertions(+), 3 deletions(-)
 create mode 100644 yarn-project/kv-store/src/lmdb/counter.test.ts

diff --git a/yarn-project/kv-store/src/lmdb/counter.test.ts b/yarn-project/kv-store/src/lmdb/counter.test.ts
new file mode 100644
index 00000000000..53ba60cdd02
--- /dev/null
+++ b/yarn-project/kv-store/src/lmdb/counter.test.ts
@@ -0,0 +1,115 @@
+import { randomBytes } from 'crypto';
+import { Database, open } from 'lmdb';
+
+import { LmdbAztecCounter } from './counter.js';
+
+describe('LmdbAztecCounter', () => {
+  let db: Database;
+
+  beforeEach(() => {
+    db = open({} as any);
+  });
+
+  describe.each([
+    ['floating point number', () => Math.random()],
+    ['integers', () => (Math.random() * 1000) | 0],
+    ['strings', () => randomBytes(8).toString('hex')],
+    ['strings', () => [Math.random(), randomBytes(8).toString('hex')]],
+  ])('counts occurrences of %s values', (_, genKey) => {
+    let counter: LmdbAztecCounter<ReturnType<typeof genKey>>;
+    beforeEach(() => {
+      counter = new LmdbAztecCounter(db, 'test');
+    });
+
+    it('returns 0 for unknown keys', () => {
+      expect(counter.get(genKey())).toEqual(0);
+    });
+
+    it('increments values', async () => {
+      const key = genKey();
+      await counter.update(key, 1);
+
+      expect(counter.get(key)).toEqual(1);
+    });
+
+    it('decrements values', async () => {
+      const key = genKey();
+      await counter.update(key, 1);
+      await counter.update(key, -1);
+
+      expect(counter.get(key)).toEqual(0);
+    });
+
+    it('increments values by a delta', async () => {
+      const key = genKey();
+      await counter.update(key, 1);
+      await counter.update(key, 2);
+
+      expect(counter.get(key)).toEqual(3);
+    });
+
+    it('resets the counter', async () => {
+      const key = genKey();
+      await counter.update(key, 1);
+      await counter.update(key, 2);
+      await counter.set(key, 0);
+
+      expect(counter.get(key)).toEqual(0);
+    });
+
+    it('iterates over entries', async () => {
+      const key = genKey();
+      await counter.update(key, 1);
+      await counter.update(key, 2);
+
+      expect([...counter.entries()]).toEqual([[key, 3]]);
+    });
+  });
+
+  it.each([
+    [
+      [
+        ['c', 2342],
+        ['a', 8],
+        ['b', 1],
+      ],
+      [
+        ['a', 8],
+        ['b', 1],
+        ['c', 2342],
+      ],
+    ],
+    [
+      [
+        [10, 2],
+        [18, 1],
+        [1, 2],
+      ],
+      [
+        [1, 2],
+        [10, 2],
+        [18, 1],
+      ],
+    ],
+    [
+      [
+        [[10, 'a'], 1],
+        [[10, 'c'], 2],
+        [[11, 'b'], 1],
+        [[9, 'f'], 1],
+        [[10, 'b'], 1],
+      ],
+      [
+        [[9, 'f'], 1],
+        [[10, 'a'], 1],
+        [[10, 'b'], 1],
+        [[10, 'c'], 2],
+        [[11, 'b'], 1],
+      ],
+    ],
+  ])('iterates in key order', async (insertOrder, expectedOrder) => {
+    const counter = new LmdbAztecCounter(db, 'test');
+    await Promise.all(insertOrder.map(([key, value]) => counter.update(key, value as number)));
+    expect([...counter.entries()]).toEqual(expectedOrder);
+  });
+});
diff --git a/yarn-project/kv-store/src/lmdb/counter.ts b/yarn-project/kv-store/src/lmdb/counter.ts
index f7b184dfdbf..4604eb74458 100644
--- a/yarn-project/kv-store/src/lmdb/counter.ts
+++ b/yarn-project/kv-store/src/lmdb/counter.ts
@@ -9,7 +9,7 @@ type CountMapKey<K> = ['count_map', string, 'slot', K];
 /**
  * A counter implementation backed by LMDB
  */
-export class LMDBCounter<K extends Key> implements AztecCounter<K> {
+export class LmdbAztecCounter<K extends Key> implements AztecCounter<K> {
   #db: Database<[K, number], CountMapKey<K>>;
   #name: string;
 
@@ -33,9 +33,16 @@ export class LMDBCounter<K extends Key> implements AztecCounter<K> {
       const slot = this.#slot(key);
       const [_, current] = this.#db.get(slot) ?? [key, 0];
       const next = current + delta;
+
+      if (next < 0) {
+        throw new Error(`Cannot update ${key} in counter ${this.#name} below zero`);
+      }
+
       if (next === 0) {
         void this.#db.remove(slot);
       } else {
+        // store the key inside the entry because LMDB might return an internal representation
+        // of the key when iterating over the database
         void this.#db.put(slot, [key, next]);
       }
 
diff --git a/yarn-project/kv-store/src/lmdb/store.ts b/yarn-project/kv-store/src/lmdb/store.ts
index 34d736244e7..8c3cebd0737 100644
--- a/yarn-project/kv-store/src/lmdb/store.ts
+++ b/yarn-project/kv-store/src/lmdb/store.ts
@@ -9,7 +9,7 @@ import { AztecMap, AztecMultiMap } from '../interfaces/map.js';
 import { AztecSingleton } from '../interfaces/singleton.js';
 import { AztecKVStore } from '../interfaces/store.js';
 import { LmdbAztecArray } from './array.js';
-import { LMDBCounter } from './counter.js';
+import { LmdbAztecCounter } from './counter.js';
 import { LmdbAztecMap } from './map.js';
 import { LmdbAztecSingleton } from './singleton.js';
 
@@ -91,7 +91,7 @@ export class AztecLmdbStore implements AztecKVStore {
   }
 
   createCounter<K extends string | number | Array<string | number>>(name: string): AztecCounter<K> {
-    return new LMDBCounter(this.#data, name);
+    return new LmdbAztecCounter(this.#data, name);
   }
 
   /**