lib/luci/net_unittest.py

# Copyright 2018 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Tests for net.py."""

import logging
import sys
import unittest

from chromite.third_party import httplib2

from chromite.lib import auth
from chromite.lib.luci import net
from chromite.lib.luci.test_support import test_case
from chromite.lib.luci.test_support import test_env


test_env.setup_test_env()


class NetTest(test_case.TestCase):
    """Tests for Net."""

    def setUp(self) -> None:
        super().setUp()

        def GetAccessToken(
            service_account_json,
        ):  # pylint: disable=unused-argument
            return "token"

        self.mock(auth, "GetAccessToken", GetAccessToken)

        self.mock(logging, "warning", lambda *_args: None)
        self.mock(logging, "error", lambda *_args: None)

    def mock_httplib2(self, calls):
        def mocked(http, **kwargs):
            if not calls:
                self.fail("Unexpected httplib2 call: %s" % kwargs)
            expected, response = calls.pop(0)
            defaults = {
                "headers": {},
                "method": "GET",
                "body": None,
            }
            defaults.update(expected)
            self.assertEqual(defaults, kwargs)
            self.assertIsInstance(http, httplib2.Http)
            if isinstance(response, Exception):
                raise response
            return response

        self.mock(net, "httprequest", mocked)
        return calls

    def test_request_works(self) -> None:
        self.mock_httplib2(
            [
                (
                    {
                        "headers": {
                            "Accept": "text/plain",
                            "Authorization": "Bearer token",
                        },
                        "method": "POST",
                        "body": "post body",
                        "uri": "http://localhost/123?a=%3D&b=%26",
                    },
                    (httplib2.Response({}), "response body"),
                ),
            ]
        )
        response = net.request(
            url="http://localhost/123",
            method="POST",
            payload="post body",
            params={"a": "=", "b": "&"},
            headers={"Accept": "text/plain"},
            include_auth=True,
            deadline=123,
            max_attempts=5,
        )
        self.assertEqual("response body", response)

    def test_retries_transient_errors(self) -> None:
        self.mock_httplib2(
            [
                ({"uri": "http://localhost/123"}, httplib2.HttpLib2Error()),
                (
                    {"uri": "http://localhost/123"},
                    (
                        httplib2.Response(
                            {
                                "status": 408,
                                "reason": "client timeout",
                            }
                        ),
                        "client timeout",
                    ),
                ),
                (
                    {"uri": "http://localhost/123"},
                    (
                        httplib2.Response(
                            {
                                "status": 500,
                                "reason": "server error",
                            }
                        ),
                        "server error",
                    ),
                ),
                (
                    {"uri": "http://localhost/123"},
                    (httplib2.Response({}), "response body"),
                ),
            ]
        )
        response = net.request("http://localhost/123", max_attempts=4)
        self.assertEqual("response body", response)

    def test_gives_up_retrying(self) -> None:
        self.mock_httplib2(
            [
                (
                    {"uri": "http://localhost/123"},
                    (
                        httplib2.Response(
                            {"status": 500, "reason": "server error"}
                        ),
                        "server error",
                    ),
                ),
                (
                    {"uri": "http://localhost/123"},
                    (
                        httplib2.Response(
                            {"status": 500, "reason": "server error"}
                        ),
                        "server error",
                    ),
                ),
                (
                    {"uri": "http://localhost/123"},
                    (httplib2.Response({}), "response body"),
                ),
            ]
        )
        with self.assertRaises(net.Error):
            net.request("http://localhost/123", max_attempts=2)

    def test_404(self) -> None:
        self.mock_httplib2(
            [
                (
                    {"uri": "http://localhost/123"},
                    (
                        httplib2.Response(
                            {"status": 404, "reason": "Not found"}
                        ),
                        "Not found",
                    ),
                ),
            ]
        )
        with self.assertRaises(net.NotFoundError):
            net.request("http://localhost/123")

    def test_crappy_cloud_endpoints_404_is_retried(self) -> None:
        self.mock_httplib2(
            [
                (
                    {"uri": "http://localhost/_ah/api/blah"},
                    (
                        httplib2.Response(
                            {"status": 404, "reason": "Not found"}
                        ),
                        "Not found",
                    ),
                ),
                (
                    {"uri": "http://localhost/_ah/api/blah"},
                    (httplib2.Response({}), "response body"),
                ),
            ]
        )
        response = net.request("http://localhost/_ah/api/blah")
        self.assertEqual("response body", response)

    def test_legitimate_cloud_endpoints_404_is_not_retried(self) -> None:
        self.mock_httplib2(
            [
                (
                    {"uri": "http://localhost/_ah/api/blah"},
                    (
                        httplib2.Response(
                            {
                                "status": 404,
                                "reason": "Not found",
                                "Content-Type": "application/json",
                            }
                        ),
                        "{}",
                    ),
                ),
            ]
        )
        with self.assertRaises(net.NotFoundError):
            net.request("http://localhost/_ah/api/blah")

    def test_401(self) -> None:
        self.mock_httplib2(
            [
                (
                    {"uri": "http://localhost/123"},
                    (
                        httplib2.Response(
                            {"status": 401, "reason": "Auth error"}
                        ),
                        "Auth error",
                    ),
                ),
            ]
        )
        with self.assertRaises(net.AuthError):
            net.request("http://localhost/123")

    def test_403(self) -> None:
        self.mock_httplib2(
            [
                (
                    {"uri": "http://localhost/123"},
                    (
                        httplib2.Response(
                            {"status": 403, "reason": "Auth error"}
                        ),
                        "Auth error",
                    ),
                ),
            ]
        )
        with self.assertRaises(net.AuthError):
            net.request("http://localhost/123")


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