swarming: add transaction_id to tasks.new request

appengine/swarming/handlers_endpoints_test.py

 #!/usr/bin/env python
 # coding=utf-8
 # Copyright 2015 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 import base64
 import datetime
 import json
 import logging
@@ skipping 45 matching lines @@
         webapp2.WSGIApplication(handlers_bot.get_routes(), debug=True),
         extra_environ={
           'REMOTE_ADDR': self.source_ip,
           'SERVER_SOFTWARE': os.environ['SERVER_SOFTWARE'],
         })
     self.mock(
         ereporter2, 'log_request',
         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
   def _enqueue_task(self, url, queue_name, **kwargs):
     if queue_name == 'task-dimensions':
       # Call directly into it, ignores any current transaction.
       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)
 
 
 class ServerApiTest(BaseTest):
   api_service_cls = handlers_endpoints.SwarmingServerService
 
   def test_details(self):
     """Asserts that server_details returns the correct version."""
     self.mock(config.config, 'config_service_hostname', lambda: 'a.server')
@@ skipping 533 matching lines @@
           ],
           u'user': u'joe@localhost',
         },
       ],
       u'now': str_now_30,
     }
     self.assertEqual(
         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)
     now = datetime.datetime(2010, 1, 2, 3, 4, 5)
     self.mock_now(now)
     str_now = unicode(now.strftime(self.DATETIME_NO_MICRO))
     request = swarming_rpcs.NewTaskRequest(
         expiration_secs=30,
         name='job1',
         priority=200,
@@ skipping 205 matching lines @@
           u'user:joe@localhost',
         ],
         u'user': u'joe@localhost',
       },
       u'task_id': u'5cee488008810',
     }
     response = self.call_api('new', body=message_to_dict(request))
     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()
     now_120_str = unicode(now_120.strftime(self.DATETIME_NO_MICRO))
 
     expected = {
       u'matched': u'2',
       u'now': now_120_str,
     }
     self.set_as_admin()
 
     def enqueue_task(*args, **kwargs):
       self.assertEqual('%s,%s' % (second, first),
                        kwargs.get('payload', ''))
       # check URL
       self.assertEqual('/internal/taskqueue/cancel-tasks', args[0])
       # check task queue
       self.assertEqual('cancel-tasks', args[1])
       return True
     self.mock(utils, 'enqueue_task', enqueue_task)
 
     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."""
     first, second, str_now_120, start, end = self._gen_two_tasks()
     first_no_perf = first.copy()
     first_no_perf.pop('performance_stats')
     # Basic request.
     request = swarming_rpcs.TasksRequest(
         end=end, start=start, include_performance_stats=True)
     expected = {u'now': str_now_120, u'items': [second, first]}
@@ skipping 165 matching lines @@
       u'ts': unicode(now.strftime(self.DATETIME_FORMAT)),
     }
     self.assertEqual(expected, self.call_api('tags', body={}).json)
 
   def _gen_two_tasks(self):
     # first request
     now = datetime.datetime(2010, 1, 2, 3, 4, 5)
     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()
     self.bot_run_task()
 
     # second request
     self.set_as_user()
     self.mock(random, 'getrandbits', lambda _: 0x88)
     now_60 = self.mock_now(now, 60)
     str_now_60 = unicode(now_60.strftime(self.DATETIME_NO_MICRO))
@@ skipping 106 matching lines @@
         utils.datetime_to_timestamp(now + datetime.timedelta(days=1)) /
         1000000.)
     self.set_as_privileged_user()
     return first, second, str_now_120, start, end
 
   def _gen_three_pending_tasks(self):
     # Creates three pending tasks, spaced 1 minute apart
     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',
         properties=dict(idempotent=True))
 
     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',
         tags=['project:yay', 'commit:efgh', 'os:Win'],
         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',
         tags=['project:yay', 'commit:ijkhl', 'os:Linux'],
         properties=dict(idempotent=True))
 
     return first_id, second_id, third_id, now, now_60, now_120
 
 
 class TaskApiTest(BaseTest):
   api_service_cls = handlers_endpoints.SwarmingTaskService
 
   def setUp(self):
     super(TaskApiTest, self).setUp()
     self.tasks_api = test_case.Endpoints(
         handlers_endpoints.SwarmingTasksService)
 
   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}
     response = self.call_api('cancel', body={'task_id': task_id})
     self.assertEqual(expected, response.json)
 
     # determine that the task's state updates correctly
     expected = {
       u'abandoned_ts': str_now,
       u'created_ts': str_now,
@@ skipping 18 matching lines @@
     # notification has been sent.
     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_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')
 
     # Attempt to cancel as non-privileged user -> HTTP 403.
     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()
     params = self.do_handshake()
     data = self.post_json('/swarming/api/v1/bot/poll', params)
     run_id = data['manifest']['task_id']
     def _params(**kwargs):
       out = {
         'cost_usd': 0.1,
         'duration': None,
@@ skipping 64 matching lines @@
     self.call_api('result', body={'task_id': '12310'}, status=404)
 
   def test_result_ok(self):
     """Asserts that result produces a result entity."""
     self.mock(random, 'getrandbits', lambda _: 0x88)
 
     # pending task
     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,
       u'failure': False,
       u'internal_failure': False,
       u'modified_ts': str_now,
       u'name': u'hi',
       u'state': u'PENDING',
       u'tags': [
         u'os:Amiga',
@@ skipping 111 matching lines @@
 
     self.set_as_privileged_user()
     run_id = task_id[:-1] + '1'
     expected = {u'output': u'rÉsult string'}
     for i in (task_id, run_id):
       response = self.call_api('stdout', body={'task_id': i})
       self.assertEqual(expected, response.json)
 
   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)
 
     run_id = task_id[:-1] + '1'
     self.call_api('stdout', body={'task_id': run_id}, status=404)
 
   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()
     self.assertEqual(task_id_1, task_id_bot[:-1] + '0')
     self.assertEqual('1', task_id_bot[-1:])
 
     # 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()
     resp = self.bot_poll()
     self.assertEqual('sleep', resp['cmd'])
 
     self.set_as_user()
 
     # results shouldn't change, even if the second task wasn't executed
     response = self.call_api('stdout', body={'task_id': task_id_2})
     self.assertEqual({'output': u'rÉsult string'}, response.json)
 
   def test_request_unknown(self):
     """Asserts that 404 is raised for unknown tasks."""
     self.call_api('request', body={'task_id': '12310'}, status=404)
 
   def test_request_ok(self):
     """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',
       u'created_ts': unicode(now.strftime(self.DATETIME_FORMAT)),
       u'expiration_secs': unicode(24 * 60 * 60),
       u'name': u'hi',
       u'priority': u'100',
       u'properties': {
         u'cipd_input': {
           u'client_package': {
@@ skipping 613 matching lines @@
     self.call_api('terminate', body={'bot_id': 'bot1'}, status=403)
 
 
 if __name__ == '__main__':
   if '-v' in sys.argv:
     unittest.TestCase.maxDiff = None
     logging.basicConfig(level=logging.DEBUG)
   else:
     logging.basicConfig(level=logging.CRITICAL)
   unittest.main()