From 5844ace43071a60b5a982f062eec305c76d5c787 Mon Sep 17 00:00:00 2001
From: Tejas Dinkar <tejas@gja.in>
Date: Thu, 3 Jan 2019 00:02:22 +0530
Subject: [PATCH] feat(KVStore): Expose the KV Store from create app (#24)

---
 app/__tests__/server_spec.js | 15 +++++++++++++++
 app/__tests__/worker_spec.js |  2 +-
 app/server.js                | 15 +++++++++++----
 app/worker.js                | 11 ++---------
 4 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/app/__tests__/server_spec.js b/app/__tests__/server_spec.js
index 8ab9c32..2fdf445 100644
--- a/app/__tests__/server_spec.js
+++ b/app/__tests__/server_spec.js
@@ -51,4 +51,19 @@ describe("server", () => {
       .get("/some-route")
       .expect(200, "127.0.0.1");
   });
+
+  it("creates stores and passes it to the worker", async () => {
+    const app = createApp(
+      'addEventListener("fetch", (e) => e.respondWith(MYSTORE.get("key").then(v => new Response(v))))',
+      {
+        kvStores: ["MYSTORE"]
+      }
+    );
+
+    await app.stores.MYSTORE.put("key", "value");
+
+    await supertest(app)
+      .get("/some-route")
+      .expect(200, "value");
+  });
 });
diff --git a/app/__tests__/worker_spec.js b/app/__tests__/worker_spec.js
index 68adc76..a4885b0 100644
--- a/app/__tests__/worker_spec.js
+++ b/app/__tests__/worker_spec.js
@@ -188,7 +188,7 @@ describe("Workers", () => {
       const worker = new Worker(
         "foo.com",
         `addEventListener("fetch", (e) => {e.respondWith(new Response("foo")); e.waitUntil(MYSTORE.put("foo", "bar"))})`,
-        { kvStores: ["MYSTORE"], kvStoreFactory: kvStoreFactory }
+        { kvStores: { MYSTORE: kvStoreFactory.getClient("MYSTORE") } }
       );
 
       const response = await worker.executeFetchEvent(`http://foo.com/blah`);
diff --git a/app/server.js b/app/server.js
index 272c002..0020e5c 100644
--- a/app/server.js
+++ b/app/server.js
@@ -21,16 +21,22 @@ async function callWorker(worker, req, res) {
   res.end(Buffer.from(data), "binary");
 }
 
-function createApp(workerContent, opts) {
+function buildKVStores(kvStoreFactory, kvStores) {
+  return kvStores.reduce((acc, name) => {
+    acc[name] = kvStoreFactory.getClient(name);
+    return acc;
+  }, {});
+}
+
+function createApp(workerContent, opts = {}) {
   let workersByOrigin = {};
-  const kvStoreFactory = new InMemoryKVStore();
+  const kvStores = buildKVStores(new InMemoryKVStore(), opts.kvStores || []);
   const app = express();
   app.use(bodyParser.raw({ type: "*/*" }));
   app.use(async (req, res) => {
     try {
       const origin = req.headers.host;
-      workersByOrigin[origin] =
-        workersByOrigin[origin] || new Worker(origin, workerContent, { kvStoreFactory, ...opts });
+      workersByOrigin[origin] = workersByOrigin[origin] || new Worker(origin, workerContent, { ...opts, kvStores });
       const worker = workersByOrigin[origin];
       await callWorker(worker, req, res);
     } catch (e) {
@@ -47,6 +53,7 @@ function createApp(workerContent, opts) {
     opts = Object.assign({}, opts, newOpts);
     workersByOrigin = {};
   };
+  app.stores = kvStores;
 
   return app;
 }
diff --git a/app/worker.js b/app/worker.js
index 89baeaf..f169558 100644
--- a/app/worker.js
+++ b/app/worker.js
@@ -7,13 +7,6 @@ const btoa = require("btoa");
 const crypto = new (require("node-webcrypto-ossl"))();
 const { TextDecoder, TextEncoder } = require("util");
 
-function buildKVStores(kvStoreFactory, kvStores) {
-  return kvStores.reduce((acc, name) => {
-    acc[name] = kvStoreFactory.getClient(name);
-    return acc;
-  }, {});
-}
-
 function chomp(str) {
   return str.substr(0, str.length - 1);
 }
@@ -40,14 +33,14 @@ function buildRequest(url, opts) {
 
 class Worker {
   constructor(origin, workerContents, opts = {}) {
-    const { upstreamHost, kvStores = [], kvStoreFactory = require("./in-memory-kv-store") } = opts;
+    const { upstreamHost, kvStores = {} } = opts;
     this.listeners = {
       fetch: e => e.respondWith(this.fetchUpstream(e.request))
     };
     this.upstreamHost = upstreamHost;
     this.origin = origin;
 
-    this.evaluateWorkerContents(workerContents, buildKVStores(kvStoreFactory, kvStores));
+    this.evaluateWorkerContents(workerContents, kvStores);
   }
 
   evaluateWorkerContents(workerContents, kvStores) {