swarming: add transaction_id to tasks.new request

appengine/swarming/handlers_endpoints_test.py

Index: appengine/swarming/handlers_endpoints_test.py
diff --git a/appengine/swarming/handlers_endpoints_test.py b/appengine/swarming/handlers_endpoints_test.py
index 5c913d139d6a6caaa96b8315a18df8ed9214b144..dbfee2445f4137672e68355aaf20cca0dd3483ec 100755
--- a/appengine/swarming/handlers_endpoints_test.py
+++ b/appengine/swarming/handlers_endpoints_test.py
@@ -63,6 +63,7 @@ class BaseTest(test_env_handlers.AppTestBase, test_case.EndpointsTestCase):
         lambda *args, **kwargs: self.fail('%s, %s' % (args, kwargs)))
     # Client API test cases run by default as user.
     self.set_as_user()
+    self.pubsub_tasks = []
     self.mock(utils, 'enqueue_task', self._enqueue_task)
 
   @ndb.non_transactional
@@ -72,6 +73,9 @@ class BaseTest(test_env_handlers.AppTestBase, test_case.EndpointsTestCase):
       handlers_backend.TaskDimensionsHandler.tidy_stale(kwargs['payload'])
       return True
     if queue_name == 'pubsub':
+      task = {'url': url, 'queue_name': queue_name}
+      task.update(kwargs)
+      self.pubsub_tasks.append(task)
       return True
     self.fail(url)
 
@@ -625,6 +629,107 @@ class TasksApiTest(BaseTest):
         expected,
         self.call_api('requests', body=message_to_dict(request)).json)
 
+  def test_new_ok_deduped_by_transaction_id(self):
+    """Asserts that new returns task result for deduped."""
+    # Run a task to completion.
+    self.mock(random, 'getrandbits', lambda _: 0x88)
+    now = self.mock_now(datetime.datetime(2010, 1, 2, 3, 4, 5))
+    str_now = unicode(now.strftime(self.DATETIME_NO_MICRO))
+    new_req, _, task_id = self.client_create_task_raw(
+        name='task', tags=['project:yay', 'commit:post'],
+        properties=dict(idempotent=True),
+        transaction_id='txn')
+    self.set_as_bot()
+    self.bot_run_task()
+
+    self.mock(random, 'getrandbits', lambda _: 0x66)
+    self.set_as_user()
+
+    expected = {
+      u'request': {
+        u'authenticated': u'user:user@example.com',
+        u'created_ts': str_now,
+        u'expiration_secs': u'86400',
+        u'name': u'task',
+        u'priority': u'100',
+        u'properties': {
+          u'cipd_input': {
+            u'client_package': {
+              u'package_name': u'infra/tools/cipd/${platform}',
+              u'version': u'git_revision:deadbeef',
+            },
+            u'packages': [{
+              u'package_name': u'rm',
+              u'path': u'bin',
+              u'version': u'git_revision:deadbeef',
+            }],
+            u'server': u'https://chrome-infra-packages.appspot.com',
+          },
+          u'command': [u'python', u'run_test.py'],
+          u'dimensions': [
+            {u'key': u'os', u'value': u'Amiga'},
+            {u'key': u'pool', u'value': u'default'},
+          ],
+          u'execution_timeout_secs': u'3600',
+          u'grace_period_secs': u'30',
+          u'idempotent': True,
+          u'io_timeout_secs': u'1200',
+          u'outputs': [u'foo', u'path/to/foobar'],
+        },
+        u'service_account': u'none',
+        u'tags': [
+          u'commit:post',
+          u'os:Amiga',
+          u'pool:default',
+          u'priority:100',
+          u'project:yay',
+          u'service_account:none',
+          u'user:joe@localhost',
+        ],
+        u'user': u'joe@localhost',
+      },
+      u'task_id': unicode(task_id),
+      u'task_result': {
+        u'bot_dimensions': [
+          {u'key': u'id', u'value': [u'bot1']},
+          {u'key': u'os', u'value': [u'Amiga']},
+          {u'key': u'pool', u'value': [u'default']},
+        ],
+        u'bot_id': u'bot1',
+        u'bot_version': self.bot_version,
+        u'completed_ts': str_now,
+        u'costs_usd': [0.1],
+        u'created_ts': str_now,
+        u'duration': 0.1,
+        u'exit_code': u'0',
+        u'failure': False,
+        u'internal_failure': False,
+        u'modified_ts': str_now,
+        u'name': u'task',
+        u'properties_hash': (
+            u'd35fe05074cbd9a2356c77c9983c71476f3a5f9415c5c5b1f3e2cdf7826b7261'
+        ),
+        u'run_id': u'5cee488008811',
+        u'server_versions': [u'v1a'],
+        u'started_ts': str_now,
+        u'state': u'COMPLETED',
+        u'tags': [
+          u'commit:post',
+          u'os:Amiga',
+          u'pool:default',
+          u'priority:100',
+          u'project:yay',
+          u'service_account:none',
+          u'user:joe@localhost',
+        ],
+        u'task_id': unicode(task_id),
+        u'try_number': u'1',
+        u'user': u'joe@localhost',
+      },
+    }
+    response = self.call_api('new', body=message_to_dict(new_req))
+    self.assertEqual(expected, response.json)
+
   def test_new_ok_isolated(self):
     """Asserts that new generates appropriate metadata."""
     self.mock(random, 'getrandbits', lambda _: 0x88)
@@ -850,12 +955,6 @@ class TasksApiTest(BaseTest):
     self.assertEqual(expected, response.json)
 
   def test_mass_cancel(self):
-    notifies = []
-    def enqueue_task_mock(**kwargs):
-      notifies.append(kwargs)
-      return True
-    self.mock(utils, 'enqueue_task', enqueue_task_mock)
-
     # Create two tasks.
     self.mock(random, 'getrandbits', lambda _: 0x88)
     first, second, _, _, _, now_120 = self._gen_three_pending_tasks()
@@ -880,7 +979,57 @@ class TasksApiTest(BaseTest):
     response = self.call_api('cancel', body={u'tags': [u'os:Win']})
     self.assertEqual(expected, response.json)
 
+  def test_cancel_by_transaction_id_ok(self):
+    """Asserts that task cancellation goes smoothly."""
+    # Create and cancel a task as a non-privileged user.
+    self.mock(random, 'getrandbits', lambda _: 0x88)
+    now = datetime.datetime(2010, 1, 2, 3, 4, 5)
+    self.mock_now(now)
+    _, _, task_id = self.client_create_task_raw(
+        pubsub_topic='projects/abc/topics/def',
+        pubsub_userdata='blah',
+        transaction_id='txn')
+    expected = {
+      u'task_id': task_id,
+      u'ok': True,
+      u'was_running': False
+    }
+    response = self.call_api(
+        'cancel_by_transaction_id', body={'transaction_id': 'txn'})
+    self.assertEqual(expected, response.json)
+
+    _, result = handlers_endpoints.get_request_and_result(task_id)
+    self.assertEqual(result.state, task_result.State.CANCELED)
+
+    expected = [
+      {
+        'payload': '{"auth_token":null,"task_id":"5cee488008810",'
+                   '"topic":"projects/abc/topics/def","userdata":"blah"}',
+        'queue_name': 'pubsub',
+        'transactional': True,
+        'url': '/internal/taskqueue/pubsub/5cee488008810',
+      },
+    ]
+    self.assertEqual(expected, self.pubsub_tasks)
+
+  def test_cancel_by_transaction_id_forbidden(self):
+    """Asserts that non-privileged non-owner can't cancel tasks."""
+    # Create a task as an admin.
+    self.mock(random, 'getrandbits', lambda _: 0x88)
+    now = datetime.datetime(2010, 1, 2, 3, 4, 5)
+    self.mock_now(now)
+    self.set_as_admin()
+    self.client_create_task_raw(
+        pubsub_topic='projects/abc/topics/def',
+        pubsub_userdata='blah',
+        transaction_id='txn')
 
+    # Attempt to cancel as non-privileged user -> HTTP 403.
+    self.set_as_user()
+    self.call_api(
+        'cancel_by_transaction_id', body={'transaction_id': 'txn'}, status=403)
+
+    self.assertEqual([], self.pubsub_tasks)
 
   def test_list_ok(self):
     """Asserts that list requests all TaskResultSummaries."""
@@ -1066,7 +1215,7 @@ class TasksApiTest(BaseTest):
     str_now = unicode(now.strftime(self.DATETIME_NO_MICRO))
     self.mock_now(now)
     self.mock(random, 'getrandbits', lambda _: 0x66)
-    _, first_id = self.client_create_task_raw(
+    _, _, first_id = self.client_create_task_raw(
         name='first', tags=['project:yay', 'commit:post', 'os:Win'],
         properties=dict(idempotent=True))
     self.set_as_bot()
@@ -1193,7 +1342,7 @@ class TasksApiTest(BaseTest):
     now = datetime.datetime(2010, 1, 2, 3, 4, 5)
     self.mock_now(now)
     self.mock(random, 'getrandbits', lambda _: 0x66)
-    _, first_id = self.client_create_task_raw(
+    _, _, first_id = self.client_create_task_raw(
         name='first', tags=['project:yay', 'commit:abcd', 'os:Win'],
         pubsub_topic='projects/abc/topics/def',
         pubsub_userdata='1234',
@@ -1201,7 +1350,7 @@ class TasksApiTest(BaseTest):
 
     now_60 = self.mock_now(now, 60)
     self.mock(random, 'getrandbits', lambda _: 0x88)
-    _, second_id = self.client_create_task_raw(
+    _, _, second_id = self.client_create_task_raw(
         name='second', user='jack@localhost',
         pubsub_topic='projects/abc/topics/def',
         pubsub_userdata='5678',
@@ -1209,7 +1358,7 @@ class TasksApiTest(BaseTest):
         properties=dict(idempotent=True))
 
     now_120 = self.mock_now(now, 120)
-    _, third_id = self.client_create_task_raw(
+    _, _, third_id = self.client_create_task_raw(
         name='third', user='jack@localhost',
         pubsub_topic='projects/abc/topics/def',
         pubsub_userdata='9000',
@@ -1229,13 +1378,12 @@ class TaskApiTest(BaseTest):
 
   def test_cancel_ok(self):
     """Asserts that task cancellation goes smoothly."""
-    # catch PubSub notification
     # Create and cancel a task as a non-privileged user.
     self.mock(random, 'getrandbits', lambda _: 0x88)
     now = datetime.datetime(2010, 1, 2, 3, 4, 5)
     self.mock_now(now)
     str_now = unicode(now.strftime(self.DATETIME_NO_MICRO))
-    _, task_id = self.client_create_task_raw(
+    _, _, task_id = self.client_create_task_raw(
         pubsub_topic='projects/abc/topics/def',
         pubsub_userdata='blah')
     expected = {u'ok': True, u'was_running': False}
@@ -1274,22 +1422,16 @@ class TaskApiTest(BaseTest):
         'url': '/internal/taskqueue/pubsub/5cee488008810',
       },
     ]
+    self.assertEqual(expected, self.pubsub_tasks)
 
   def test_cancel_forbidden(self):
     """Asserts that non-privileged non-owner can't cancel tasks."""
-    # catch PubSub notification
-    notifies = []
-    def enqueue_task_mock(**kwargs):
-      notifies.append(kwargs)
-      return True
-    self.mock(utils, 'enqueue_task', enqueue_task_mock)
-
     # Create a task as an admin.
     self.mock(random, 'getrandbits', lambda _: 0x88)
     now = datetime.datetime(2010, 1, 2, 3, 4, 5)
     self.mock_now(now)
     self.set_as_admin()
-    _, task_id = self.client_create_task_raw(
+    _, _, task_id = self.client_create_task_raw(
         pubsub_topic='projects/abc/topics/def',
         pubsub_userdata='blah')
 
@@ -1297,12 +1439,14 @@ class TaskApiTest(BaseTest):
     self.set_as_user()
     self.call_api('cancel', body={'task_id': task_id}, status=403)
 
+    self.assertEqual([], self.pubsub_tasks)
+
   def test_task_canceled(self):
     self.mock(random, 'getrandbits', lambda _: 0x88)
     now = datetime.datetime(2010, 1, 2, 3, 4, 5)
     self.mock_now(now)
     str_now = unicode(now.strftime(self.DATETIME_NO_MICRO))
-    _, task_id = self.client_create_task_raw(
+    _, _, task_id = self.client_create_task_raw(
         properties=dict(command=['python', 'runtest.py']))
 
     self.set_as_bot()
@@ -1387,7 +1531,7 @@ class TaskApiTest(BaseTest):
     now = datetime.datetime(2010, 1, 2, 3, 4, 5)
     self.mock_now(now)
     str_now = unicode(now.strftime(self.DATETIME_NO_MICRO))
-    _, task_id = self.client_create_task_raw()
+    _, _, task_id = self.client_create_task_raw()
     response = self.call_api('result', body={'task_id': task_id})
     expected = {
       u'created_ts': str_now,
@@ -1519,7 +1663,7 @@ class TaskApiTest(BaseTest):
 
   def test_stdout_empty(self):
     """Asserts that incipient tasks produce no output."""
-    _, task_id = self.client_create_task_raw()
+    _, _, task_id = self.client_create_task_raw()
     response = self.call_api('stdout', body={'task_id': task_id})
     self.assertEqual({}, response.json)
 
@@ -1528,13 +1672,14 @@ class TaskApiTest(BaseTest):
 
   def test_result_run_not_found(self):
     """Asserts that getting results from incipient tasks raises 404."""
-    _, task_id = self.client_create_task_raw()
+    _, _, task_id = self.client_create_task_raw()
     run_id = task_id[:-1] + '1'
     self.call_api('stdout', body={'task_id': run_id}, status=404)
 
   def test_task_deduped(self):
     """Asserts that task deduplication works as expected."""
-    _, task_id_1 = self.client_create_task_raw(properties=dict(idempotent=True))
+    _, _, task_id_1 = self.client_create_task_raw(
+        properties=dict(idempotent=True))
 
     self.set_as_bot()
     task_id_bot = self.bot_run_task()
@@ -1543,7 +1688,7 @@ class TaskApiTest(BaseTest):
 
     # second task; this one's results should be returned immediately
     self.set_as_user()
-    _, task_id_2 = self.client_create_task_raw(
+    _, _, task_id_2 = self.client_create_task_raw(
         name='second', user='jack@localhost', properties=dict(idempotent=True))
 
     self.set_as_bot()
@@ -1564,7 +1709,7 @@ class TaskApiTest(BaseTest):
     """Asserts that request produces a task request."""
     now = datetime.datetime(2010, 1, 2, 3, 4, 5, 6)
     self.mock_now(now)
-    _, task_id = self.client_create_task_raw()
+    _, _, task_id = self.client_create_task_raw()
     response = self.call_api('request', body={'task_id': task_id})
     expected = {
       u'authenticated': u'user:user@example.com',