appengine/swarming/server/resultdb_test.py

#!/usr/bin/env vpython
# Copyright 2020 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 datetime
import logging
import sys
import unittest
import uuid

import mock

import test_env

test_env.setup_test_env()

from google.appengine.api import app_identity
from google.appengine.ext import ndb

from test_support import test_case

from components import net

from server import config
from server import resultdb

from proto.config import config_pb2


class ResultDBTest(test_case.TestCase):
  # This test needs to be run independently.
  # run by test.py
  no_run = 1

  def setUp(self):
    super(ResultDBTest, self).setUp()
    mock.patch(
      "google.appengine.api.app_identity.get_default_version_hostname"
    ).start().return_value = "test-swarming.appspot.com"
    mock.patch("uuid.uuid4").start().return_value = uuid.UUID(int=0)

  def tearDown(self):
    super(ResultDBTest, self).tearDown()
    mock.patch.stopall()

  @ndb.tasklet
  def nop_async(self, *_args, **_kwargs):
    pass

  @ndb.tasklet
  def _mock_call_resultdb_recorder_api_async(
    self, _method, _request, response_headers, **_kwargs
  ):
    response_headers["update-token"] = "token"

  def test_create_invocation_async(self):

    with mock.patch(
      "server.resultdb._call_resultdb_recorder_api_async",
      mock.MagicMock(side_effect=self._mock_call_resultdb_recorder_api_async),
    ) as mock_call:
      update_token = resultdb.create_invocation_async(
        "task001", "infra:try", datetime.datetime(2020, 1, 2, 3, 4, 5)
      ).get_result()
      self.assertEqual(update_token, "token")
      mock_call.assert_called_once_with(
        "CreateInvocation",
        {
          "invocation": {
            "producerResource": "//test-swarming.appspot.com/tasks/task001",
            "realm": "infra:try",
            "deadline": "2020-01-02T03:04:05Z",
          },
          "requestId": "00000000-0000-0000-0000-000000000000",
          "invocationId": "task-test-swarming.appspot.com-task001",
        },
        project_id="infra",
        response_headers=mock.ANY,
      )

  def test_create_invocation_async_no_update_token(self):
    with mock.patch(
      "server.resultdb._call_resultdb_recorder_api_async",
      mock.Mock(side_effect=self.nop_async),
    ):
      with self.assertRaisesRegexp(
        AssertionError, "^response_headers should have valid update-token: {}$"
      ):
        resultdb.create_invocation_async(
          "task001",
          "infra:try",
          datetime.datetime(2020, 1, 2, 3, 4, 5),
        ).get_result()

  def test_finalize_invocation_async_success(self):

    with mock.patch(
      "server.resultdb._call_resultdb_recorder_api_async"
    ) as mock_call:
      mock_call.side_effect = self.nop_async
      resultdb.finalize_invocation_async("task001", "secret").get_result()
      mock_call.assert_called_once_with(
        "FinalizeInvocation",
        {
          "name": "invocations/task-test-swarming.appspot.com-task001",
        },
        headers={"update-token": "secret"},
        scopes=None,
      )

  def test_finalize_invocation_async_failed_precondition(self):
    with mock.patch(
      "server.resultdb._call_resultdb_recorder_api_async"
    ) as mock_call:
      mock_call.side_effect = net.Error(
        msg="error",
        status_code=400,
        response="error",
        headers={"X-Prpc-Grpc-Code": "9"},
      )
      resultdb.finalize_invocation_async("task001", "secret").get_result()
      mock_call.assert_called_once_with(
        "FinalizeInvocation",
        {
          "name": "invocations/task-test-swarming.appspot.com-task001",
        },
        headers={"update-token": "secret"},
        scopes=None,
      )

  def test_finalize_invocation_async_failed(self):
    with mock.patch(
      "server.resultdb._call_resultdb_recorder_api_async"
    ) as mock_call:
      mock_call.side_effect = Exception("failed")
      with self.assertRaises(Exception):
        resultdb.finalize_invocation_async("task001", "secret").get_result()
      mock_call.assert_called_once_with(
        "FinalizeInvocation",
        {
          "name": "invocations/task-test-swarming.appspot.com-task001",
        },
        headers={"update-token": "secret"},
        scopes=None,
      )

  def test_call_resultdb_recorder_api(self):

    with (
      mock.patch("server.config.settings") as mock_settings,
      mock.patch(
        "components.net.json_request_async"
      ) as mock_json_request_async,
    ):
      mock_settings.return_value = config_pb2.SettingsCfg(
        resultdb=config_pb2.ResultDBSettings(
          server="https://results.api.cr.dev"
        ),
      )
      mock_json_request_async.side_effect = self.nop_async

      resultdb._call_resultdb_recorder_api_async(
        "FinalizeInvocation", {}
      ).get_result()

      mock_settings.assert_called_once()
      mock_json_request_async.assert_called_once()

  def test_get_invocation_name(self):
    self.assertEqual(
      resultdb.get_invocation_name("run_result_id_01"),
      "invocations/task-test-swarming.appspot.com-run_result_id_01",
    )


if __name__ == "__main__":
  if "-v" in sys.argv:
    unittest.TestCase.maxDiff = None
  logging.basicConfig(
    level=logging.DEBUG if "-v" in sys.argv else logging.CRITICAL
  )
  unittest.main()