From 1eda868ea234b543e138fd0b1a21238bf97e35fc Mon Sep 17 00:00:00 2001
From: Matt Carroll <matthewcarroll@google.com>
Date: Fri, 28 Sep 2018 14:58:25 -0700
Subject: [PATCH] Update context API

Change-Id: I2ac7672fead481c6d97846fb5735a9672761b5cb
---
 .eslintrc.json                |   7 +-
 .gitignore                    |   3 +-
 README.md                     |   7 -
 generateDocs.js               |   2 +-
 package.json                  |   3 +-
 src/contexts.js               | 268 ++++++++++++++++++++++++++++++++
 src/dialogflow-fulfillment.js |  60 ++++---
 src/v1-agent.js               |  19 ++-
 src/v2-agent.js               |  26 ++--
 test/contexts-test.js         | 283 ++++++++++++++++++++++++++++++++++
 test/webhook-v1-test.js       |  14 +-
 test/webhook-v2-test.js       |  28 ++--
 12 files changed, 646 insertions(+), 74 deletions(-)
 create mode 100644 src/contexts.js
 create mode 100644 test/contexts-test.js

diff --git a/.eslintrc.json b/.eslintrc.json
index db67580..b6af202 100644
--- a/.eslintrc.json
+++ b/.eslintrc.json
@@ -1,9 +1,12 @@
 {
     "extends": "google",
     "parserOptions": {
-	   "ecmaVersion": 8
+       "ecmaVersion": 8,
+       "ecmaFeatures": {
+            "experimentalObjectRestSpread": true
+        }
     },
     "rules": {
         "max-len": ["error", 120]
     }
-}
\ No newline at end of file
+}
diff --git a/.gitignore b/.gitignore
index a3c7cf8..50a85c7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,4 +6,5 @@ coverage
 *.tgz
 package-lock.json
 webhook-client.md
-rich-responses.md
\ No newline at end of file
+rich-responses.md
+.nyc_output
\ No newline at end of file
diff --git a/README.md b/README.md
index be1674f..cf53e5e 100644
--- a/README.md
+++ b/README.md
@@ -1,10 +1,8 @@
 # Dialogflow Fulfillment Library
-
 The [Dialogflow Fulfillment Library](https://dialogflow.com/docs/fulfillment) allows you to connect natural language understanding and processing to your own systems, APIs, and databases. Using Fulfillment, you can surface commands and information from your services to your users through a natural conversational interface.
 
 Dialogflow Fulfillment makes creating fulfillment for Dialogflow v1 and v2 agents for 8 chat and voice platforms on Node.js easy and simple.
 
-
 ![fulfillment library works with 8 platforms](https://raw.githubusercontent.com/dialogflow/dialogflow-fulfillment-nodejs/master/dialogflow-fulfillment-graphic.png "Dialogflow's fulfillment library works with 8 platforms")
 
 ## Supported features
@@ -18,18 +16,14 @@ This library is intended to help build Node.js Dialogflow Fulfillment for multip
 
 If only building Dialogflow Fulfillment for the [Google Assistant](https://dialogflow.com/docs/integrations/google-assistant) and no other integrations, use the Actions of Google NPM module ([actions-on-google](https://www.npmjs.com/package/actions-on-google)) which supports all Actions on Google features.
 
-
 ## Quick Start
-
 1. [Sign-up/Log-in to Dialogflow](https://console.dialogflow.com/api-client/#/login)
 2. Create a Dialogflow agent
 3. Go to **Fulfillment** > **Enable Dialogflow Inline Editor**<sup> A.</sup> > **package.json** tab to add `"dialogflow-fulfillment": "^0.5.0"` to the `dependencies` object.
 4. Select **Deploy**.
-
   <sup>A.</sup> Powered by Cloud Functions for Firebase
 
 ## Setup Instructions
-
 ```javascript
 // Import the appropriate class
 const { WebhookClient } = require('dialogflow-fulfillment');
@@ -37,7 +31,6 @@ const { WebhookClient } = require('dialogflow-fulfillment');
  //Create an instance
 const agent = new WebhookClient({request: request, response: response});
 ```
-
 ## Samples
 | Name                                 | Language                         |
 | ------------------------------------ |:---------------------------------|
diff --git a/generateDocs.js b/generateDocs.js
index 99b0798..abd2df3 100644
--- a/generateDocs.js
+++ b/generateDocs.js
@@ -24,7 +24,7 @@ const inputFiles = ['./src/*.js', './src/rich-responses/*.js'];
 const outputDir = './docs/';
 
 const webhookFilename = 'webhook-client.md';
-const webhookClientClassNames = ['WebhookClient', 'V2Agent', 'V1Agent'];
+const webhookClientClassNames = ['WebhookClient', 'V2Agent', 'V1Agent', 'Context'];
 const richResponseFilename = 'rich-responses.md'
 const richRepsonseClassNames = [ 'RichResponse', 'Card', 'Suggestion', 'Image', 'Payload', 'Text'];
 
diff --git a/package.json b/package.json
index 509ddb8..c156718 100644
--- a/package.json
+++ b/package.json
@@ -37,7 +37,8 @@
     "fulfillment"
   ],
   "dependencies": {
-    "debug": "^3.1.0"
+    "debug": "^3.1.0",
+    "lodash": "^4.17.11"
   },
   "peerDependencies": {
     "actions-on-google": "^2.1.3"
diff --git a/src/contexts.js b/src/contexts.js
new file mode 100644
index 0000000..f0ebbe3
--- /dev/null
+++ b/src/contexts.js
@@ -0,0 +1,268 @@
+/**
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+const debug = require('debug')('dialogflow:debug');
+const _ = require('lodash');
+
+// Configure logging for hosting platforms that only support console.log and console.error
+debug.log = console.log.bind(console);
+
+const DELETED_LIFESPAN_COUNT = 0; // Lifespan of a deleted context
+
+
+/**
+ * This is the class that handles Dialogflow's contexts for the WebhookClient class
+ */
+class Context {
+  /**
+   * Constructor for Context object
+   * To be used in by Dialogflow's webhook client class
+   * context objects take are formatted as follows:
+   *   { "context-name": {
+   *        "lifespan": 5,
+   *        "parameters": {
+   *          "param": "value"
+   *        }
+   *     }
+   *   }
+   *
+   * @example
+   * const context = new Context(inputContexts);
+   * context.get('name of context')
+   * context.set('another context name', 5, {param: 'value'})
+   * context.delete('name of context') // set context lifespan to 0
+   *
+   * @param {Object} inputContexts input contexts of a v1 or v2 webhook request
+   * @param {string} session for a v2 webhook request & response
+   */
+  constructor(inputContexts, session) {
+    /**
+     * Dialogflow contexts included in the request or empty object if no value
+     * https://dialogflow.com/docs/contexts
+     * @type {object}
+     */
+    this.contexts = {};
+    this.session = session;
+    if (inputContexts && session) {
+      this.inputContexts = this._processV2InputContexts(inputContexts);
+      this.contexts = this._processV2InputContexts(inputContexts);
+    } else if (inputContexts) {
+      this.contexts = this._processV1InputContexts(inputContexts);
+      this.inputContexts = this._processV1InputContexts(inputContexts);
+    }
+  }
+  // ---------------------------------------------------------------------------
+  //              Public CRUD methods
+  // ---------------------------------------------------------------------------
+  /**
+   * Set a new Dialogflow outgoing context: https://dialogflow.com/docs/contexts
+   *
+   * @example
+   * const { WebhookClient } = require('dialogflow-webhook');
+   * const agent = new WebhookClient({request: request, response: response});
+   * agent.context.set('sample context name');
+   * const context = {'name': 'weather', 'lifespan': 2, 'parameters': {'city': 'Rome'}};
+   *
+   * @param {string|Object} name of context or an object representing a context
+   * @param {number} [lifespan=5] lifespan of context, number with a value of 0 or greater
+   * @param {Object} [params] parameters of context (can be arbitrary key-value pairs)
+   */
+  set(name, lifespan, params) {
+    if (!name || (typeof name !== 'string' && typeof name['name'] !== 'string')) {
+      throw new Error('Required "name" argument must be a string or an object with a string attribute "name"');
+    }
+    if (typeof name !== 'string') {
+      params = name['parameters'];
+      lifespan = name['lifespan'];
+      name = name['name'];
+    }
+    if (!this.contexts[name]) {
+      this.contexts[name] = {name: name};
+    }
+    if (lifespan !== undefined && lifespan !== null) {
+      this.contexts[name].lifespan = lifespan;
+    }
+    if (params !== undefined) {
+      this.contexts[name].parameters = params;
+    }
+  }
+
+  /**
+   * Get an context from the Dialogflow webhook request: https://dialogflow.com/docs/contexts
+   *
+   * @example
+   * const { WebhookClient } = require('dialogflow-webhook');
+   * const agent = new WebhookClient({request: request, response: response});
+   * let context = agent.context.get('sample context name');
+   *
+   * @param {string} name of an context present in the Dialogflow webhook request
+   * @return {Object|null} context object with lifespan and parameters (if defined) or null
+   */
+  get(name) {
+    return this.contexts[name];
+  }
+  /**
+   * Delete an context a Dialogflow session (set the lifespan to 0)
+   *
+   * @example
+   * const { WebhookClient } = require('dialogflow-webhook');
+   * const agent = new WebhookClient({request: request, response: response});
+   * agent.context.delete('no-longer-relevant-context-name');
+   *
+   * @param {string} name of context to be deleted
+   *
+   * @public
+   */
+  delete(name) {
+    this.set(name, DELETED_LIFESPAN_COUNT);
+  }
+  /**
+   * Returns contexts as an iterator.
+   *
+   * @example
+   * const { WebhookClient } = require('dialogflow-webhook');
+   * const agent = new WebhookClient({request: request, response: response});
+   * for (const context of agent.context) {
+   *   // do something with the contexts
+   * }
+   *
+   * @return {iterator} iterator of all context objects
+   * @public
+   */
+  [Symbol.iterator]() {
+    let contextArray = [];
+    for (const contextName of Object.keys(this.contexts)) {
+      contextArray.push(this.contexts[contextName]);
+    }
+    return contextArray[Symbol.iterator]();
+    // suppose to be Array.prototype.values(), but can't use because of bug:
+    // https://bugs.chromium.org/p/chromium/issues/detail?id=615873
+  }
+  // ---------------------------------------------------------------------------
+  //              Private methods
+  // ---------------------------------------------------------------------------
+  /**
+   * Remove an context from Dialogflow's outgoing context webhook response
+   * used to maintain compatibility with legacy clearContext methods
+   *
+   * @example
+   * const { WebhookClient } = require('dialogflow-webhook');
+   * const agent = new WebhookClient({request: request, response: response});
+   * agent.context._removeOutgoingContext('no-longer-sent-context-name');
+   *
+   * @param {string} name of context to be removed from outgoing contexts
+   *
+   * @private
+   */
+  _removeOutgoingContext(name) {
+    delete this.contexts[name];
+  }
+  // ---------------------------------------------------------------------------
+  //              Private v2 <--> v1 translation methods
+  // ---------------------------------------------------------------------------
+  /**
+   * Translate context object from v1 webhook request format to class format
+   *
+   * @param {Array} v1InputContexts to be used by the Contexts class
+   *
+   * @return {Object} internal representation of contexts
+   * @private
+   */
+  _processV1InputContexts(v1InputContexts) {
+    let contexts = {};
+    for (let index = 0; index<v1InputContexts.length; index++) {
+      const context = v1InputContexts[index];
+      contexts[context['name']] = {
+        name: context['name'],
+        parameters: context['parameters'],
+        lifespan: context['lifespan'],
+      };
+    }
+    return contexts;
+  }
+  /**
+   * Translate context object from v2 webhook request format to class format
+   *
+   * @param {Array} v2InputContexts to be used by the Contexts class
+   *
+   * @return {Object} internal representation of contexts
+   * @private
+   */
+  _processV2InputContexts(v2InputContexts) {
+    let contexts = {};
+    for (let index = 0; index<v2InputContexts.length; index++) {
+      let context = v2InputContexts[index];
+      const name = context['name'].split('/')[6];
+      contexts[name] = {
+        name: name,
+        lifespan: context['lifespanCount'],
+        parameters: context['parameters']};
+    }
+    return contexts;
+  }
+  /**
+   * Get array of context objects formatted for v1 webhook response
+   *
+   * @return {Object[]} array of v1 context objects for webhook response
+   */
+  getV1OutputContextsArray() {
+    let v1OutputContexts = [];
+    for (const ctx of this) {
+      // Skip context if it is the same as the input context
+      if (this.inputContexts &&
+        this.inputContexts[ctx.name] &&
+        _.isEqual(ctx, this.inputContexts[ctx.name])) {
+        continue;
+      }
+      let v1Context = {name: ctx.name};
+      if (ctx.lifespan !== undefined) {
+        v1Context['lifespan'] = ctx.lifespan;
+      }
+      if (ctx.parameters) {
+        v1Context['parameters'] = ctx.parameters;
+      }
+      v1OutputContexts.push(v1Context);
+    }
+    return v1OutputContexts;
+  }
+  /**
+   * Get array of context objects formatted for v2 webhook response
+   *
+   * @return {Object[]} array of v2 context objects for webhook response
+   */
+  getV2OutputContextsArray() {
+    let v2OutputContexts = [];
+    for (const ctx of this) {
+      // Skip context if it is the same as the input context
+      if (this.inputContexts &&
+        this.inputContexts[ctx.name] &&
+        _.isEqual(ctx, this.inputContexts[ctx.name])) {
+        continue;
+      }
+      let v2Context = {name: `${this.session}/contexts/${ctx.name}`};
+      if (ctx.lifespan !== undefined) {
+        v2Context['lifespanCount'] = ctx.lifespan;
+      }
+      if (ctx.parameters) {
+        v2Context['parameters'] = ctx.parameters;
+      }
+      v2OutputContexts.push(v2Context);
+    }
+    return v2OutputContexts;
+  }
+}
+
+module.exports = Context;
diff --git a/src/dialogflow-fulfillment.js b/src/dialogflow-fulfillment.js
index 2c8aef8..5ab7957 100644
--- a/src/dialogflow-fulfillment.js
+++ b/src/dialogflow-fulfillment.js
@@ -144,9 +144,17 @@ class WebhookClient {
      * Dialogflow contexts included in the request or null if no value
      * https://dialogflow.com/docs/contexts
      * @type {string}
+     * @deprecated
      */
     this.contexts = null;
 
+    /**
+     * Instance of Dialogflow contexts class to provide an API to set/get/delete contexts
+     *
+     * @type {Contexts}
+     */
+    this.context = null;
+
     /**
      * Dialogflow source included in the request or null if no value
      * https://dialogflow.com/docs/reference/agent/query#query_parameters_and_json_fields
@@ -321,7 +329,7 @@ class WebhookClient {
   }
 
   // --------------------------------------------------------------------------
-  //          Context and follow-up event methods
+  //          Deprecated Context methods
   // --------------------------------------------------------------------------
   /**
    * Set a new Dialogflow outgoing context: https://dialogflow.com/docs/contexts
@@ -335,18 +343,15 @@ class WebhookClient {
    *
    * @param {string|Object} context name of context or an object representing a context
    * @return {WebhookClient}
+   * @deprecated
    */
   setContext(context) {
-    // If developer provides a string, transform to context object, using string as the name
+    console.warn('setContext is deprecated, migrate to `contexts.set`');
     if (typeof context === 'string') {
-      context = {name: context};
-    }
-    if (context && !context.name) {
-      throw new Error('context must be provided and must have a name');
+      this.context.set(context);
+    } else {
+      this.context.set(context.name, context.lifespan, context.parameters);
     }
-
-    this.client.addContext_(context);
-
     return this;
   }
 
@@ -359,9 +364,13 @@ class WebhookClient {
    * agent.clearOutgoingContexts();
    *
    * @return {WebhookClient}
+   * @deprecated
    */
   clearOutgoingContexts() {
-    this.outgoingContexts_ = [];
+    console.warn('clearOutgoingContexts is deprecated, migrate to `contexts.delete` or `contexts.set`');
+    for (const ctx of this.context) {
+      this.context._removeOutgoingContext(ctx.name);
+    }
     return this;
   }
 
@@ -375,20 +384,11 @@ class WebhookClient {
    *
    * @param {string} context name of an existing outgoing context
    * @return {WebhookClient}
+   * @deprecated
    */
   clearContext(context) {
-    if (this.agentVersion === 1) {
-      this.outgoingContexts_ = this.outgoingContexts_.filter(
-        (ctx) => ctx.name !== context
-      );
-    } else if (this.agentVersion === 2) {
-      // Take all existing outgoing contexts and filter out the context that needs to be cleared
-      this.outgoingContexts_ = this.outgoingContexts_.filter(
-        (ctx) => ctx.name.slice(-context.length) !== context
-      );
-    } else {
-      debug('Couldn\'t find context');
-    }
+    console.warn('clearContext is deprecated, migrate to `contexts.delete` or `contexts.set`');
+    this.context._removeOutgoingContext(context);
     return this;
   }
 
@@ -402,11 +402,25 @@ class WebhookClient {
    *
    * @param {string} contextName name of an context present in the Dialogflow webhook request
    * @return {Object} context context object with the context name
+   * @deprecated
    */
   getContext(contextName) {
-    return this.contexts.filter( (context) => context.name === contextName )[0] || null;
+    console.warn('getContext is deprecated, migrate to `contexts.get`');
+    const context = this.context.get(contextName);
+    if (context) {
+      return {
+        name: contextName,
+        lifespan: context.lifespan,
+        parameters: context.parameters,
+      };
+    } else {
+      return null;
+    }
   }
 
+  // --------------------------------------------------------------------------
+  //          Follow-up event method
+  // --------------------------------------------------------------------------
   /**
    * Set the followup event
    *
diff --git a/src/v1-agent.js b/src/v1-agent.js
index cca8afa..14ef445 100644
--- a/src/v1-agent.js
+++ b/src/v1-agent.js
@@ -30,6 +30,9 @@ const Image = require('./rich-responses/image-response');
 const Suggestion = require('./rich-responses/suggestions-response');
 const PayloadResponse = require('./rich-responses/payload-response');
 
+// Contexts class
+const Contexts = require('./contexts');
+
 /**
  * Class representing a v1 Dialogflow agent
  */
@@ -84,6 +87,13 @@ class V1Agent {
     this.agent.contexts = this.agent.request_.body.result.contexts;
     debug(`Input contexts: ${JSON.stringify(this.agent.contexts)}`);
 
+    /**
+     * Instance of Dialogflow contexts class to provide an API to set/get/delete contexts
+     *
+     * @type {Contexts}
+     */
+    this.agent.context = new Contexts(this.agent.request_.body.result.contexts);
+
     /**
      * Dialogflow source included in the request or null if no value
      * https://dialogflow.com/docs/reference/agent/query#query_parameters_and_json_fields
@@ -203,7 +213,7 @@ class V1Agent {
       throw new Error(`No responses defined for platform: ${requestSource}`);
     }
 
-    responseJson.contextOut = this.agent.outgoingContexts_;
+    responseJson.contextOut = this.agent.context.getV1OutputContextsArray();
     this.agent.followupEvent_ ? responseJson.followupEvent = this.agent.followupEvent_ : undefined;
 
     debug('Response to Dialogflow: ' + JSON.stringify(responseJson));
@@ -234,7 +244,7 @@ class V1Agent {
    */
   addContext_(context) {
     // v1 contexts have the same structure as used by the library
-    this.agent.outgoingContexts_.push(context);
+    this.agent.context.set(context);
   }
 
   /**
@@ -247,8 +257,9 @@ class V1Agent {
     let eventJson = {
       name: event.name,
     };
-    event.parameters ? eventJson.data = event.parameters : undefined;
-
+    if (event.parameters) {
+      eventJson.data = event.parameters;
+    }
     this.agent.followupEvent_ = eventJson;
   }
 
diff --git a/src/v2-agent.js b/src/v2-agent.js
index 6df11e2..8efbe26 100644
--- a/src/v2-agent.js
+++ b/src/v2-agent.js
@@ -19,8 +19,6 @@ const debug = require('debug')('dialogflow:debug');
 // Configure logging for hosting platforms agent only support console.log and console.error
 debug.log = console.log.bind(console);
 
-const DEFAULT_CONTEXT_LIFESPAN = 5;
-
 // Response Builder classes
 const {
   V1_TO_V2_PLATFORM_NAME,
@@ -32,6 +30,9 @@ const Image = require('./rich-responses/image-response');
 const Suggestion = require('./rich-responses/suggestions-response');
 const PayloadResponse = require('./rich-responses/payload-response');
 
+// Contexts class
+const Contexts = require('./contexts');
+
 /**
  * Class representing a v2 Dialogflow agent
  */
@@ -102,6 +103,15 @@ class V2Agent {
     }
     debug(`Request contexts: ${JSON.stringify(this.agent.contexts)}`);
 
+    /**
+     * Instance of Dialogflow contexts class to provide an API to set/get/delete contexts
+     *
+     * @type {Contexts}
+     */
+    this.agent.context = new Contexts(
+      this.agent.request_.body.queryResult.outputContexts,
+      this.agent.session);
+
     /**
      * Dialogflow source included in the request or null if no value
      * https://dialogflow.com/docs/reference/api-v2/rest/v2beta1/projects.agent.intents#Platform
@@ -233,7 +243,7 @@ class V2Agent {
       throw new Error(`No responses defined for platform: ${requestSource}`);
     }
 
-    responseJson.outputContexts = this.agent.outgoingContexts_;
+    responseJson.outputContexts = this.agent.context.getV2OutputContextsArray();
     if (this.agent.followupEvent_) {
       responseJson.followupEventInput = this.agent.followupEvent_;
     }
@@ -270,15 +280,7 @@ class V2Agent {
     if (context.name.match('/contexts/')) {
       context = this.convertV2ContextToV1Context_(context);
     }
-
-    // v2 contexts require the use of the session name and a transformation
-    // from a v1 context object to a v2 context object before adding
-    let v2Context = {};
-    v2Context.name = this.agent.session + '/contexts/' + context.name;
-    v2Context.lifespanCount = context.lifespan === undefined ? DEFAULT_CONTEXT_LIFESPAN: context.lifespan;
-    v2Context.parameters = context.parameters;
-
-    this.agent.outgoingContexts_.push(v2Context);
+    this.agent.context.set(context);
   }
 
   /**
diff --git a/test/contexts-test.js b/test/contexts-test.js
new file mode 100644
index 0000000..b426492
--- /dev/null
+++ b/test/contexts-test.js
@@ -0,0 +1,283 @@
+/**
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the 'License');
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+'use strict';
+
+// Enable dialogflow debug logging
+// process.env.DEBUG = 'dialogflow:*';
+
+const test = require('ava');
+
+const Context = require('../src/contexts');
+
+test('Test Context constructor', async (t) => {
+  // TextResponse generic response
+  let context = new Context();
+  t.deepEqual(context.contexts, {});
+
+  context = new Context(v1IncomingContexts);
+  t.deepEqual(context.contexts, contextObjects);
+
+  context = new Context(v2IncomingContexts, v2Session);
+  t.deepEqual(context.contexts, contextObjects);
+});
+
+test('Test v1& v2 JSON output', async (t) => {
+  // TextResponse generic response
+  let context = new Context();
+  t.deepEqual(context.getV1OutputContextsArray(), []);
+  t.deepEqual(context.getV2OutputContextsArray(), []);
+
+  context = new Context(v1IncomingContexts);
+  t.deepEqual(context.getV1OutputContextsArray(), []);
+
+  context = new Context();
+  context.set(contextObject1);
+  context.set(contextObject2);
+  context.set(contextObject3);
+  t.deepEqual(context.getV1OutputContextsArray(), v1IncomingContexts);
+
+  context = new Context(v2IncomingContexts, v2Session);
+  t.deepEqual(context.getV2OutputContextsArray(), []);
+
+  context = new Context(null, v2Session);
+  context.set(contextObject1);
+  context.set(contextObject2);
+  context.set(contextObject3);
+  t.deepEqual(context.getV2OutputContextsArray(), v2IncomingContexts);
+});
+
+test('Test set method', async (t) => {
+  // Set contexts using context objects
+  let context = new Context();
+  context.set(contextObject1);
+  t.deepEqual(context.contexts['context name'], contextObject1);
+  context.set(contextObject2);
+  context.set(contextObject3);
+  t.deepEqual(context.contexts, contextObjects);
+
+  // Set contexts using method arguments
+  context = new Context();
+  context.set('context name', 0, {});
+  t.deepEqual(context.contexts['context name'], contextObject1);
+  context.set('another context name', 99, {parameter: 'value'});
+  context.set(
+    'yet another context name',
+    4,
+    {parameter: 'value', anotherParam: 'another value', yetAnotherParam: 'yet another value'}
+  );
+  t.deepEqual(context.contexts, contextObjects);
+
+  // Overwriting existing context values with context objects
+  context = new Context();
+  context.set(contextObject1);
+  t.deepEqual(context.contexts['context name'], contextObject1);
+  // overwrite lifespan
+  contextObject1.lifespan = 45;
+  context.set(contextObject1);
+  t.deepEqual(context.contexts['context name'], contextObject1);
+  // overwrite parameters
+  contextObject1.parameters = {param: 'value'};
+  context.set(contextObject1);
+  t.deepEqual(context.contexts['context name'], contextObject1);
+
+  // Overwriting existing context values with method arguments
+  context = new Context();
+  context.set(contextObject2);
+  t.deepEqual(context.contexts['another context name'], contextObject2);
+  // overwrite lifespan
+  contextObject2.lifespan = 45;
+  context.set('another context name', 45);
+  t.deepEqual(context.contexts['another context name'], contextObject2);
+  // overwrite parameters
+  contextObject2.parameters = {param: 'value'};
+  context.set('another context name', null, {param: 'value'});
+  t.deepEqual(context.contexts['another context name'], contextObject2);
+
+  // Cleanup
+  contextObject1 = {name: 'context name', lifespan: 0, parameters: {}};
+  contextObject2 = {name: 'another context name', parameters: {parameter: 'value'}, lifespan: 99};
+});
+
+test('Test error for set method required argument', async (t) => {
+  let context = new Context();
+  const noNameDefinedError = t.throws(() => {
+   context.set();
+  });
+  t.is(
+    noNameDefinedError.message,
+    'Required "name" argument must be a string or an object with a string attribute "name"'
+  );
+});
+
+test('Test get method', async (t) => {
+  // Get contexts set with constructor
+  let context = new Context(v2IncomingContexts, v2Session);
+  t.deepEqual(context.get('context name'), contextObject1);
+  t.deepEqual(context.get('another context name'), contextObject2);
+  t.deepEqual(context.get('yet another context name'), contextObject3);
+
+  // Get contexts set with set (context object)
+  context = new Context();
+  context.set(contextObject3);
+  t.deepEqual(context.get('yet another context name'), contextObject3);
+
+  // Get contexts set with set (method parameters)
+  context = new Context();
+  context.set(
+    'yet another context name',
+    4,
+    {parameter: 'value', anotherParam: 'another value', yetAnotherParam: 'yet another value'}
+  );
+  t.deepEqual(context.get('yet another context name'), contextObject3);
+});
+
+test('Test delete method', async (t) => {
+  // delete with incoming contexts
+  let context = new Context(v2IncomingContexts, v2Session);
+  context.delete('another context name');
+  t.deepEqual(context.contexts['another context name'].lifespan, 0);
+  t.deepEqual(context.contexts['yet another context name'].lifespan, 4);
+
+  // delete with contexts set via context object
+  context = new Context();
+  context.set(contextObject2);
+  context.delete('another context name');
+  t.deepEqual(context.contexts['another context name'].lifespan, 0);
+
+  // delete with contexts set via method parameters
+  context = new Context();
+  context.set('another context name', 45, null);
+  context.delete('another context name');
+  t.deepEqual(context.contexts['another context name'].lifespan, 0);
+});
+
+test('Test context iterator', async (t) => {
+  let context = new Context(v2IncomingContexts, v2Session);
+  let i = 0;
+  for (const ctx of context) {
+    t.deepEqual(ctx, [contextObject1, contextObject2, contextObject3][i]);
+    i++;
+  }
+});
+
+test('Test _removeOutgoingContext method', async (t) => {
+  let context = new Context(v2IncomingContexts, v2Session);
+  context._removeOutgoingContext('another context name');
+  t.deepEqual(context.contexts, {'context name': contextObject1, 'yet another context name': contextObject3});
+});
+
+
+// ---------------------------------------------------------------------------
+//              Context test helper objects
+// ---------------------------------------------------------------------------
+let contextObject1 = {
+  name: 'context name',
+  lifespan: 0,
+  parameters: {},
+};
+
+let contextObject2 = {
+  name: 'another context name',
+  parameters: {
+    parameter: 'value',
+  },
+  lifespan: 99,
+};
+
+let contextObject3 = {
+  name: 'yet another context name',
+  parameters: {
+    parameter: 'value',
+    anotherParam: 'another value',
+    yetAnotherParam: 'yet another value',
+  },
+  lifespan: 4,
+};
+
+const v1IncomingContexts = [
+  {
+    name: 'context name',
+    parameters: {},
+    lifespan: 0,
+  },
+  {
+    name: 'another context name',
+    parameters: {
+      parameter: 'value',
+    },
+    lifespan: 99,
+  },
+  {
+    name: 'yet another context name',
+    parameters: {
+      parameter: 'value',
+      anotherParam: 'another value',
+      yetAnotherParam: 'yet another value',
+    },
+    lifespan: 4,
+  },
+];
+
+const v2Session = 'projects/project-id/agent/sessions/88d1...a0';
+
+const v2IncomingContexts = [
+  {
+    name: 'projects/project-id/agent/sessions/88d1...a0/contexts/context name',
+    parameters: {},
+    lifespanCount: 0,
+  },
+  {
+    name: 'projects/project-id/agent/sessions/88d1...a0/contexts/another context name',
+    parameters: {
+      parameter: 'value',
+    },
+    lifespanCount: 99,
+  },
+  {
+    name: 'projects/project-id/agent/sessions/88d1...a0/contexts/yet another context name',
+    parameters: {
+      parameter: 'value',
+      anotherParam: 'another value',
+      yetAnotherParam: 'yet another value',
+    },
+    lifespanCount: 4,
+  },
+];
+
+const contextObjects = {
+  'context name': {
+    name: 'context name',
+    lifespan: 0,
+    parameters: {},
+  },
+  'another context name': {
+    name: 'another context name',
+    parameters: {
+      parameter: 'value',
+    },
+    lifespan: 99,
+  },
+  'yet another context name': {
+    name: 'yet another context name',
+    parameters: {
+      parameter: 'value',
+      anotherParam: 'another value',
+      yetAnotherParam: 'yet another value',
+    },
+    lifespan: 4,
+  },
+};
diff --git a/test/webhook-v1-test.js b/test/webhook-v1-test.js
index 5ec8d20..46609a5 100644
--- a/test/webhook-v1-test.js
+++ b/test/webhook-v1-test.js
@@ -197,17 +197,21 @@ test('Test v1 contexts', async (t) => {
   });
   // setContext
   agent.setContext(sampleContextName);
-  t.deepEqual({name: sampleContextName}, agent.outgoingContexts_[0]);
+  t.deepEqual({name: sampleContextName}, agent.context.get(sampleContextName));
   agent.setContext(secondContextName);
-  t.deepEqual({name: secondContextName}, agent.outgoingContexts_[1]);
+  t.deepEqual({name: secondContextName}, agent.context.get(secondContextName));
   agent.setContext(complexContext);
-  t.deepEqual(complexContext, agent.outgoingContexts_[2]);
+  t.deepEqual({name: complexContext.name,
+    lifespan: 2, parameters: {city: 'Rome'}},
+    agent.context.get(complexContext.name)
+  );
   // clearContext
   agent.clearContext(sampleContextName);
-  t.deepEqual({name: secondContextName}, agent.outgoingContexts_[0]);
+  t.deepEqual(undefined, agent.context.get(sampleContextName));
   // clearAllContext
   agent.clearOutgoingContexts();
-  t.deepEqual([], agent.outgoingContexts_);
+  t.deepEqual([], agent.context.getV1OutputContextsArray());
+  t.deepEqual([], agent.context.getV2OutputContextsArray());
 });
 
 test('Test v1 getContext', async (t) => {
diff --git a/test/webhook-v2-test.js b/test/webhook-v2-test.js
index 3a9941c..26ac90f 100644
--- a/test/webhook-v2-test.js
+++ b/test/webhook-v2-test.js
@@ -258,23 +258,12 @@ test('Test v2 Twitter payload response', async (t) => {
 test('Test v2 contexts', async (t) => {
   const v2Request = mockGoogleV2Request;
   const sampleContextName = 'sample context name';
-  const sampleV2ContextName =
-    v2Request.session + '/contexts/' + sampleContextName;
   const secondContextName = 'second context name';
-  const secondV2ContextName =
-    v2Request.session + '/contexts/' + secondContextName;
   const complexContext = {
     name: 'weather',
     lifespan: 2,
     parameters: {city: 'Rome'},
   };
-  const complexV2ContextName =
-    v2Request.session + '/contexts/' + complexContext.name;
-  const complexV2Context = {
-    name: complexV2ContextName,
-    lifespanCount: 2,
-    parameters: {city: 'Rome'},
-  };
 
   let googleResponse = new ResponseMock();
   let googleRequest = {body: v2Request};
@@ -285,21 +274,24 @@ test('Test v2 contexts', async (t) => {
   // setContext
   agent.setContext(sampleContextName);
   t.deepEqual(
-    {name: sampleV2ContextName, lifespanCount: 5, parameters: undefined},
-    agent.outgoingContexts_[0]
+    {name: sampleContextName},
+    agent.context.get(sampleContextName)
   );
   agent.setContext(secondContextName);
   t.deepEqual(
-    {name: secondV2ContextName, lifespanCount: 5, parameters: undefined},
-    agent.outgoingContexts_[1]
+    {name: secondContextName},
+    agent.context.get(secondContextName)
   );
   agent.setContext(complexContext);
-  t.deepEqual(complexV2Context, agent.outgoingContexts_[2]);
+  t.deepEqual(
+    {name: complexContext.name, lifespan: 2, parameters: {city: 'Rome'}},
+    agent.context.get(complexContext.name)
+  );
   // clearContext
   agent.clearContext(sampleContextName);
   t.deepEqual(
-    {name: secondV2ContextName, lifespanCount: 5, parameters: undefined},
-    agent.outgoingContexts_[0]
+    undefined,
+    agent.context.get(sampleContextName)
   );
   // clearAllContext
   agent.clearOutgoingContexts();