fix(toolkit): correctly reset context from the shell command (#1903)

packages/aws-cdk/lib/commands/context.ts

@@ -2,7 +2,7 @@ import colors = require('colors/safe');
 import yargs = require('yargs');
 import { CommandOptions } from '../command-api';
 import { print } from '../logging';
-import { PROJECT_CONFIG } from '../settings';
+import { Context, PROJECT_CONFIG } from '../settings';
 import { renderTable } from '../util';
 
 export const command = 'context';
@@ -34,8 +34,7 @@ export async function realHandler(options: CommandOptions): Promise<number> {
     await configuration.saveContext();
     print('All context values cleared.');
   } else if (args.reset) {
-    invalidateContext(contextValues, args.reset);
-    configuration.context.setAll(contextValues);
+    invalidateContext(configuration.context, args.reset);
     await configuration.saveContext();
   } else {
     // List -- support '--json' flag
@@ -77,16 +76,16 @@ function listContext(context: any) {
   print(`Run ${colors.blue('cdk context --reset KEY_OR_NUMBER')} to remove a context key. It will be refreshed on the next CDK synthesis run.`);
 }
 
-function invalidateContext(context: any, key: string) {
+function invalidateContext(context: Context, key: string) {
   const i = parseInt(key, 10);
   if (`${i}` === key) {
     // Twas a number and we fully parsed it.
     key = keyByNumber(context, i);
   }
 
   // Unset!
-  if (key in context) {
-    delete context[key];
+  if (context.has(key)) {
+    context.unset(key);
     print(`Context value ${colors.blue(key)} reset. It will be refreshed on the next SDK synthesis run.`);
   } else {
     print(`No context value with key ${colors.blue(key)}`);

packages/aws-cdk/lib/settings.ts

@@ -24,6 +24,7 @@ export class Configuration {
   private readonly commandLineArguments: Settings;
   private projectConfig: Settings;
   private projectContext: Settings;
+  private loaded = false;
 
   constructor(commandLineArguments?: yargs.Arguments) {
     this.commandLineArguments = commandLineArguments
@@ -48,13 +49,17 @@ export class Configuration {
       .merge(this.commandLineArguments)
       .makeReadOnly();
 
+    this.loaded = true;
+
     return this;
   }
 
   /**
    * Save the project config
    */
   public async saveContext(): Promise<this> {
+    if (!this.loaded) { return this; }
+
     if (this.context.modifiedBottom) {
       await this.projectConfig.save(PROJECT_CONFIG);
     }
@@ -85,6 +90,14 @@ export class Context {
   constructor(private readonly bottom: Settings, private readonly bottomPrefixPath: string[], private readonly top: Settings) {
   }
 
+  public get keys(): string[] {
+    return Object.keys(this.everything());
+  }
+
+  public has(key: string) {
+    return this.keys.indexOf(key) > -1;
+  }
+
   public everything(): {[key: string]: any} {
     const b = this.bottom.get(this.bottomPrefixPath) || {};
     const t = this.top.get([]) || {};
@@ -105,18 +118,12 @@ export class Context {
     }
   }
 
-  public setAll(values: object) {
-    for (const [key, value] of Object.entries(values)) {
-      this.set(key, value);
-    }
-  }
-
   public unset(key: string) {
     this.set(key, undefined);
   }
 
   public clear() {
-    for (const key of Object.keys(this.everything())) {
+    for (const key of this.keys) {
       this.unset(key);
     }
   }

packages/aws-cdk/test/commands/test.context-command.ts

@@ -0,0 +1,25 @@
+import { Test } from 'nodeunit';
+import { realHandler } from '../../lib/commands/context';
+import { Configuration } from '../../lib/settings';
+
+export = {
+  async 'context reset can remove a context key'(test: Test) {
+    // GIVEN
+    const configuration = new Configuration();
+    configuration.context.set('foo', 'bar');
+    configuration.context.set('baz', 'quux');
+
+    // WHEN
+    await realHandler({
+      configuration,
+      args: { reset: 'foo' }
+    } as any);
+
+    // THEN
+    test.deepEqual(configuration.context.everything(), {
+      baz: 'quux'
+    });
+
+    test.done();
+  },
+};