tests/swarming_test.py

#!/usr/bin/env vpython
# -*- coding: utf-8 -*-
# Copyright 2013 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.

from __future__ import print_function

import logging
import os
import StringIO
import sys
import traceback

# Mutates sys.path.
import test_env

# third_party/
from depot_tools import auto_stub

import net_utils

import auth
import swarming
from utils import logging_utils


def main(args):
  """Bypasses swarming.main()'s exception handling.

  It gets in the way when debugging test failures.
  """
  dispatcher = swarming.subcommand.CommandDispatcher("swarming")
  return dispatcher.execute(swarming.OptionParserSwarming(), args)


class Common:
  # pylint: disable=no-member

  def setUp(self):
    self.mock(auth, "ensure_logged_in", lambda _: None)
    self.mock(sys, "stdout", StringIO.StringIO())
    self.mock(sys, "stderr", StringIO.StringIO())
    self.mock(logging_utils, "prepare_logging", lambda *args: None)
    self.mock(logging_utils, "set_console_level", lambda *args: None)

  def tearDown(self):
    if not self.has_failed():
      self._check_output("", "")

  maxDiff = None

  def _check_output(self, out, err):
    self.assertMultiLineEqual(out, sys.stdout.getvalue())
    self.assertMultiLineEqual(err, sys.stderr.getvalue())

    # Flush their content by mocking them again.
    self.mock(sys, "stdout", StringIO.StringIO())
    self.mock(sys, "stderr", StringIO.StringIO())

  def main_safe(self, args):
    """Bypasses swarming.main()'s exception handling.

    It gets in the way when debugging test failures.
    """
    # pylint: disable=bare-except
    try:
      return main(args)
    except:
      data = "%s\nSTDOUT:\n%s\nSTDERR:\n%s" % (
        traceback.format_exc(),
        sys.stdout.getvalue(),
        sys.stderr.getvalue(),
      )
      self.fail(data)


class NetTestCase(net_utils.TestCase, Common):
  # These test fail when running in parallel
  # Need to run in sequential_test_runner.py as an executable
  no_run = 1
  """Base class that defines the url_open mock."""

  def setUp(self):
    net_utils.TestCase.setUp(self)
    Common.setUp(self)


class TestMain(NetTestCase):
  # Tests calling main().
  def test_query_base(self):
    self.expected_requests(
      [
        (
          "https://localhost:1/_ah/api/swarming/v1/bot/botid/tasks?limit=200",
          {},
          {"yo": "dawg"},
        ),
      ]
    )
    ret = self.main_safe(
      [
        "query",
        "--swarming",
        "https://localhost:1",
        "bot/botid/tasks",
      ]
    )
    self._check_output('{\n  "yo": "dawg"\n}\n', "")
    self.assertEqual(0, ret)

  def test_query_cursor(self):
    self.expected_requests(
      [
        (
          "https://localhost:1/_ah/api/swarming/v1/bot/botid/tasks?"
          "foo=bar&limit=2",
          {},
          {
            "cursor": "%",
            "extra": False,
            "items": ["A"],
          },
        ),
        (
          "https://localhost:1/_ah/api/swarming/v1/bot/botid/tasks?"
          "foo=bar&cursor=%25&limit=1",
          {},
          {
            "cursor": None,
            "items": ["B"],
            "ignored": True,
          },
        ),
      ]
    )
    ret = self.main_safe(
      [
        "query",
        "--swarming",
        "https://localhost:1",
        "bot/botid/tasks?foo=bar",
        "--limit",
        "2",
      ]
    )
    expected = (
      '{\n  "extra": false, \n  "items": [\n    "A", \n    "B"\n  ]\n}\n'
    )
    self._check_output(expected, "")
    self.assertEqual(0, ret)


if __name__ == "__main__":
  for env_var_to_remove in ("SWARMING_TASK_ID", "SWARMING_SERVER"):
    os.environ.pop(env_var_to_remove, None)
  test_env.main()