lib/luci/prpc/client_unittest.py

# -*- coding: utf-8 -*-
# Copyright 2018 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# pylint: disable=unused-argument

from __future__ import print_function

import contextlib
import sys
import unittest

from google.protobuf import empty_pb2
import mock

from chromite.lib.luci import net
from chromite.lib.luci.prpc import client as prpc_client
from chromite.lib.luci.prpc import codes
from chromite.lib.luci.prpc.test import test_pb2, test_prpc_pb2
from chromite.lib.luci.test_support import test_case, test_env

test_env.setup_test_env()


class PRPCClientTestCase(test_case.TestCase):

  def make_test_client(self):
    return prpc_client.Client('localhost', test_prpc_pb2.TestServiceDescription)

  def test_generated_methods(self):
    expected_methods = {
        'Give',
        'Take',
        'Echo',
    }

    members = dir(self.make_test_client())
    for m in expected_methods:
      self.assertIn(m, members)

  @contextlib.contextmanager
  def mocked_request(self, res=None):
    res = res or empty_pb2.Empty()
    with mock.patch('chromite.lib.luci.net.request', autospec=True) as m:
      m.return_value = res.SerializeToString()
      yield

  def test_request(self):
    with self.mocked_request():
      req = test_pb2.GiveRequest(m=1)
      self.make_test_client().Give(req)

      net.request.assert_called_with(
          url='https://localhost/prpc/test.Test/Give',
          method='POST',
          payload=req.SerializeToString(),
          headers={
              'Content-Type': 'application/prpc; encoding=binary',
              'Accept': 'application/prpc; encoding=binary',
              'X-Prpc-Timeout': '10S',
          },
          include_auth=True,
          deadline=10,
          max_attempts=4,
      )

  def give_creds(self, creds):
    self.make_test_client().Give(test_pb2.GiveRequest(), credentials=creds)

  def test_request_credentials_service_account(self):
    with self.mocked_request():
      self.give_creds(prpc_client.include_auth())
      _, kwargs = net.request.call_args
      self.assertEqual(kwargs['include_auth'], True)

  def test_request_credentials_service_account_key(self):
    with self.mocked_request():
      self.give_creds(prpc_client.include_auth())
      _, kwargs = net.request.call_args
      self.assertEqual(kwargs['include_auth'], True)

  def test_request_timeout(self):
    with self.mocked_request():
      self.make_test_client().Give(test_pb2.GiveRequest(), timeout=20)
      _, kwargs = net.request.call_args
      self.assertEqual(kwargs['deadline'], 20)
      self.assertEqual(kwargs['headers']['X-Prpc-Timeout'], '20S')

  def test_response_ok(self):
    expected = test_pb2.TakeResponse(k=1)
    with self.mocked_request(res=expected):
      actual = self.make_test_client().Take(empty_pb2.Empty())
      self.assertEqual(actual, expected)

  @mock.patch('chromite.lib.luci.net.request', autospec=True)
  def test_response_protocol_error(self, request):
    request.side_effect = net.NotFoundError(
        msg='not found',
        status_code=404,
        response='not found',
        headers={
            # no X-Prpc-Grpc-Code header.
        },
    )
    with self.assertRaises(prpc_client.ProtocolError):
      self.make_test_client().Take(empty_pb2.Empty())

  @mock.patch('chromite.lib.luci.net.request', autospec=True)
  def test_response_rpc_error(self, request):
    request.side_effect = net.NotFoundError(
        msg='not found',
        status_code=404,
        response=b'not found',
        headers={
            'X-Prpc-Grpc-Code': str(codes.StatusCode.NOT_FOUND[0]),
        },
    )
    with self.assertRaises(prpc_client.RpcError) as cm:
      self.make_test_client().Take(empty_pb2.Empty())
    self.assertEqual(cm.exception.status_code, codes.StatusCode.NOT_FOUND)


if __name__ == '__main__':
  if '-v' in sys.argv:
    unittest.TestCase.maxDiff = None
  unittest.main()