Allows both --json and token args on create/update (#930)

twosigma · Aug 22, 2019 · cd5eada · cd5eada
1 parent 4b449dd
commit cd5eada
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 10 deletions.
diff --git a/cli/integration/tests/waiter/test_cli.py b/cli/integration/tests/waiter/test_cli.py
@@ -809,9 +809,35 @@ def test_update_token_json(self):
 
     def test_post_token_json_and_flags(self):
         token_name = self.token_name()
-        cp = cli.update(self.waiter_url, token_name, update_flags='--json foo.json --cpus 0.1')
-        self.assertEqual(1, cp.returncode, cp.stderr)
-        self.assertIn('cannot specify both a token JSON file and token field flags', cli.stderr(cp))
+        update_fields = {'cpus': 0.2, 'mem': 256}
+        with cli.temp_token_file(update_fields) as path:
+            cp = cli.update(self.waiter_url, token_name, update_flags=f'--json {path} --cpus 0.1')
+            self.assertEqual(1, cp.returncode, cp.stderr)
+            self.assertIn('cannot specify the same parameter in both a token JSON file and token field flags at the '
+                          'same time (cpus)', cli.stderr(cp))
+
+            cp = cli.update(self.waiter_url, token_name, update_flags=f'--json {path} --mem 128')
+            self.assertEqual(1, cp.returncode, cp.stderr)
+            self.assertIn('cannot specify the same parameter in both a token JSON file and token field flags at the '
+                          'same time (mem)', cli.stderr(cp))
+
+            cp = cli.update(self.waiter_url, token_name, update_flags=f'--json {path} --cpus 0.1 --mem 128')
+            self.assertEqual(1, cp.returncode, cp.stderr)
+            self.assertIn('cannot specify the same parameter in both a token JSON file and token field flags',
+                          cli.stderr(cp))
+            self.assertIn('cpus', cli.stderr(cp))
+            self.assertIn('mem', cli.stderr(cp))
+
+            try:
+                cp = cli.update(self.waiter_url, token_name, update_flags=f'--json {path} --version foo --image bar')
+                self.assertEqual(0, cp.returncode, cp.stderr)
+                token_data = util.load_token(self.waiter_url, token_name)
+                self.assertEqual(0.2, token_data['cpus'])
+                self.assertEqual(256, token_data['mem'])
+                self.assertEqual('foo', token_data['version'])
+                self.assertEqual('bar', token_data['image'])
+            finally:
+                util.delete_token(self.waiter_url, token_name)
 
     def test_post_token_json_invalid(self):
         token_name = self.token_name()

diff --git a/cli/waiter/token_post.py b/cli/waiter/token_post.py
@@ -100,19 +100,23 @@ def create_or_update_token(clusters, args, _, action):
     token_name_from_args = args.pop('token', None)
     json_file = args.pop('json', None)
     if json_file:
-        if len(args) > 0:
-            raise Exception('You cannot specify both a token JSON file and token field flags at the same time.')
-
         if json_file == '-':
             token_json = read_token_from_stdin()
-            token_fields = parse_raw_token_spec(token_json)
+            token_fields_from_json = parse_raw_token_spec(token_json)
         else:
-            token_fields = load_json_file(json_file)
-            if not token_fields:
+            token_fields_from_json = load_json_file(json_file)
+            if not token_fields_from_json:
                 raise Exception(f'Unable to load token JSON from {json_file}.')
     else:
-        token_fields = args
+        token_fields_from_json = {}
+
+    token_fields_from_args = args
+    shared_keys = set(token_fields_from_json).intersection(token_fields_from_args)
+    if shared_keys:
+        raise Exception(f'You cannot specify the same parameter in both a token JSON file and token field flags at '
+                        f'the same time ({", ".join(shared_keys)}).')
 
+    token_fields = {**token_fields_from_json, **token_fields_from_args}
     token_name_from_json = token_fields.pop('token', None)
     if token_name_from_args and token_name_from_json:
         raise Exception('You cannot specify the token name both as an argument and in the token JSON file.')