# Copyright (C) The Arvados Authors. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

import errno
import hashlib
import mmap
import os
import random
import re
import shutil
import socket
import stat
import sys
import tempfile
import time
import unittest
import urllib.parse

from pathlib import Path
from unittest import mock
from unittest.mock import patch

import parameterized
import pycurl

import arvados
import arvados.retry
import arvados.util
from . import arvados_testutil as tutil
from . import keepstub
from . import run_test_server

from .arvados_testutil import DiskCacheBase

@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepTestCase(run_test_server.TestCaseWithServers, DiskCacheBase):
    disk_cache = False
    MAIN_SERVER = {}
    KEEP_SERVER = {}
    block_cache_test = None

    @classmethod
    def setUpClass(cls):
        super(KeepTestCase, cls).setUpClass()
        run_test_server.authorize_with("admin")
        cls.api_client = arvados.api('v1')
        cls.block_cache_test = DiskCacheBase()
        cls.keep_client = arvados.KeepClient(api_client=cls.api_client,
                                             proxy='', local_store='',
                                             block_cache=cls.block_cache_test.make_block_cache(cls.disk_cache))

    @classmethod
    def tearDownClass(cls):
        super(KeepTestCase, cls).setUpClass()
        cls.block_cache_test.tearDown()

    def test_KeepBasicRWTest(self):
        self.assertEqual(0, self.keep_client.upload_counter.get())
        foo_locator = self.keep_client.put('foo')
        self.assertRegex(
            foo_locator,
            '^acbd18db4cc2f85cedef654fccc4a4d8\+3',
            'wrong md5 hash from Keep.put("foo"): ' + foo_locator)

        # 6 bytes because uploaded 2 copies
        self.assertEqual(6, self.keep_client.upload_counter.get())

        self.assertEqual(0, self.keep_client.download_counter.get())
        self.assertTrue(tutil.binary_compare(self.keep_client.get(foo_locator),
                         b'foo'),
                         'wrong content from Keep.get(md5("foo"))')
        self.assertEqual(3, self.keep_client.download_counter.get())

    def test_KeepBinaryRWTest(self):
        blob_str = b'\xff\xfe\xf7\x00\x01\x02'
        blob_locator = self.keep_client.put(blob_str)
        self.assertRegex(
            blob_locator,
            '^7fc7c53b45e53926ba52821140fef396\+6',
            ('wrong locator from Keep.put(<binarydata>):' + blob_locator))
        self.assertEqual(self.keep_client.get(blob_locator),
                         blob_str,
                         'wrong content from Keep.get(md5(<binarydata>))')

    def test_KeepLongBinaryRWTest(self):
        blob_data = b'\xff\xfe\xfd\xfc\x00\x01\x02\x03'
        for i in range(0, 23):
            blob_data = blob_data + blob_data
        blob_locator = self.keep_client.put(blob_data)
        self.assertRegex(
            blob_locator,
            '^84d90fc0d8175dd5dcfab04b999bc956\+67108864',
            ('wrong locator from Keep.put(<binarydata>): ' + blob_locator))
        self.assertEqual(self.keep_client.get(blob_locator),
                         blob_data,
                         'wrong content from Keep.get(md5(<binarydata>))')

    @unittest.skip("unreliable test - please fix and close #8752")
    def test_KeepSingleCopyRWTest(self):
        blob_data = b'\xff\xfe\xfd\xfc\x00\x01\x02\x03'
        blob_locator = self.keep_client.put(blob_data, copies=1)
        self.assertRegex(
            blob_locator,
            '^c902006bc98a3eb4a3663b65ab4a6fab\+8',
            ('wrong locator from Keep.put(<binarydata>): ' + blob_locator))
        self.assertEqual(self.keep_client.get(blob_locator),
                         blob_data,
                         'wrong content from Keep.get(md5(<binarydata>))')

    def test_KeepEmptyCollectionTest(self):
        blob_locator = self.keep_client.put('', copies=1)
        self.assertRegex(
            blob_locator,
            '^d41d8cd98f00b204e9800998ecf8427e\+0',
            ('wrong locator from Keep.put(""): ' + blob_locator))

    def test_unicode_must_be_ascii(self):
        # If unicode type, must only consist of valid ASCII
        foo_locator = self.keep_client.put(u'foo')
        self.assertRegex(
            foo_locator,
            '^acbd18db4cc2f85cedef654fccc4a4d8\+3',
            'wrong md5 hash from Keep.put("foo"): ' + foo_locator)

        if sys.version_info < (3, 0):
            with self.assertRaises(UnicodeEncodeError):
                # Error if it is not ASCII
                self.keep_client.put(u'\xe2')

        with self.assertRaises(AttributeError):
            # Must be bytes or have an encode() method
            self.keep_client.put({})

    def test_KeepHeadTest(self):
        locator = self.keep_client.put('test_head')
        self.assertRegex(
            locator,
            '^b9a772c7049325feb7130fff1f8333e9\+9',
            'wrong md5 hash from Keep.put for "test_head": ' + locator)
        self.assertEqual(True, self.keep_client.head(locator))
        self.assertEqual(self.keep_client.get(locator),
                         b'test_head',
                         'wrong content from Keep.get for "test_head"')

@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepPermissionTestCase(run_test_server.TestCaseWithServers, DiskCacheBase):
    disk_cache = False
    MAIN_SERVER = {}
    KEEP_SERVER = {'blob_signing': True}

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def test_KeepBasicRWTest(self):
        run_test_server.authorize_with('active')
        keep_client = arvados.KeepClient(block_cache=self.make_block_cache(self.disk_cache))
        foo_locator = keep_client.put('foo')
        self.assertRegex(
            foo_locator,
            r'^acbd18db4cc2f85cedef654fccc4a4d8\+3\+A[a-f0-9]+@[a-f0-9]+$',
            'invalid locator from Keep.put("foo"): ' + foo_locator)
        self.assertEqual(keep_client.get(foo_locator),
                         b'foo',
                         'wrong content from Keep.get(md5("foo"))')

        # GET with an unsigned locator => bad request
        bar_locator = keep_client.put('bar')
        unsigned_bar_locator = "37b51d194a7513e45b56f6524f2d51f2+3"
        self.assertRegex(
            bar_locator,
            r'^37b51d194a7513e45b56f6524f2d51f2\+3\+A[a-f0-9]+@[a-f0-9]+$',
            'invalid locator from Keep.put("bar"): ' + bar_locator)
        self.assertRaises(arvados.errors.KeepReadError,
                          keep_client.get,
                          unsigned_bar_locator)

        # GET from a different user => bad request
        run_test_server.authorize_with('spectator')
        self.assertRaises(arvados.errors.KeepReadError,
                          arvados.Keep.get,
                          bar_locator)

        # Unauthenticated GET for a signed locator => bad request
        # Unauthenticated GET for an unsigned locator => bad request
        keep_client.api_token = ''
        self.assertRaises(arvados.errors.KeepReadError,
                          keep_client.get,
                          bar_locator)
        self.assertRaises(arvados.errors.KeepReadError,
                          keep_client.get,
                          unsigned_bar_locator)


@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepProxyTestCase(run_test_server.TestCaseWithServers, DiskCacheBase):
    disk_cache = False
    MAIN_SERVER = {}
    KEEP_SERVER = {}
    KEEP_PROXY_SERVER = {}

    @classmethod
    def setUpClass(cls):
        super(KeepProxyTestCase, cls).setUpClass()
        run_test_server.authorize_with('active')
        cls.api_client = arvados.api('v1')

    def tearDown(self):
        super(KeepProxyTestCase, self).tearDown()
        DiskCacheBase.tearDown(self)

    def test_KeepProxyTest1(self):
        # Will use ARVADOS_KEEP_SERVICES environment variable that
        # is set by setUpClass().
        keep_client = arvados.KeepClient(api_client=self.api_client,
                                         local_store='', block_cache=self.make_block_cache(self.disk_cache))
        baz_locator = keep_client.put('baz')
        self.assertRegex(
            baz_locator,
            '^73feffa4b7f6bb68e44cf984c85f6e88\+3',
            'wrong md5 hash from Keep.put("baz"): ' + baz_locator)
        self.assertEqual(keep_client.get(baz_locator),
                         b'baz',
                         'wrong content from Keep.get(md5("baz"))')
        self.assertTrue(keep_client.using_proxy)

    def test_KeepProxyTestMultipleURIs(self):
        # Test using ARVADOS_KEEP_SERVICES env var overriding any
        # existing proxy setting and setting multiple proxies
        arvados.config.settings()['ARVADOS_KEEP_SERVICES'] = 'http://10.0.0.1 https://foo.example.org:1234/'
        keep_client = arvados.KeepClient(api_client=self.api_client,
                                         local_store='',
                                         block_cache=self.make_block_cache(self.disk_cache))
        uris = [x['_service_root'] for x in keep_client._keep_services]
        self.assertEqual(uris, ['http://10.0.0.1/',
                                'https://foo.example.org:1234/'])

    def test_KeepProxyTestInvalidURI(self):
        arvados.config.settings()['ARVADOS_KEEP_SERVICES'] = 'bad.uri.org'
        with self.assertRaises(arvados.errors.ArgumentError):
            keep_client = arvados.KeepClient(api_client=self.api_client,
                                             local_store='',
                                             block_cache=self.make_block_cache(self.disk_cache))


@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepClientServiceTestCase(unittest.TestCase, tutil.ApiClientMock, DiskCacheBase):
    disk_cache = False

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def get_service_roots(self, api_client):
        keep_client = arvados.KeepClient(api_client=api_client, block_cache=self.make_block_cache(self.disk_cache))
        services = keep_client.weighted_service_roots(arvados.KeepLocator('0'*32))
        return [urllib.parse.urlparse(url) for url in sorted(services)]

    def test_ssl_flag_respected_in_roots(self):
        for ssl_flag in [False, True]:
            services = self.get_service_roots(self.mock_keep_services(
                service_ssl_flag=ssl_flag))
            self.assertEqual(
                ('https' if ssl_flag else 'http'), services[0].scheme)

    def test_correct_ports_with_ipv6_addresses(self):
        service = self.get_service_roots(self.mock_keep_services(
            service_type='proxy', service_host='100::1', service_port=10, count=1))[0]
        self.assertEqual('100::1', service.hostname)
        self.assertEqual(10, service.port)

    def test_recognize_proxy_services_in_controller_response(self):
        keep_client = arvados.KeepClient(api_client=self.mock_keep_services(
            service_type='proxy', service_host='localhost', service_port=9, count=1),
                                         block_cache=self.make_block_cache(self.disk_cache))
        try:
            # this will fail, but it ensures we get the service
            # discovery response
            keep_client.put('baz2', num_retries=0)
        except:
            pass
        self.assertTrue(keep_client.using_proxy)

    def test_insecure_disables_tls_verify(self):
        api_client = self.mock_keep_services(count=1)
        force_timeout = socket.timeout("timed out")

        api_client.insecure = True
        with tutil.mock_keep_responses(b'foo', 200) as mock:
            keep_client = arvados.KeepClient(api_client=api_client, block_cache=self.make_block_cache(self.disk_cache))
            keep_client.get('acbd18db4cc2f85cedef654fccc4a4d8+3')
            self.assertEqual(
                mock.responses[0].getopt(pycurl.SSL_VERIFYPEER),
                0)
            self.assertEqual(
                mock.responses[0].getopt(pycurl.SSL_VERIFYHOST),
                0)

        api_client.insecure = False
        with tutil.mock_keep_responses(b'foo', 200) as mock:
            keep_client = arvados.KeepClient(api_client=api_client, block_cache=self.make_block_cache(self.disk_cache))
            keep_client.get('acbd18db4cc2f85cedef654fccc4a4d8+3')
            # getopt()==None here means we didn't change the
            # default. If we were using real pycurl instead of a mock,
            # it would return the default value 1.
            self.assertEqual(
                mock.responses[0].getopt(pycurl.SSL_VERIFYPEER),
                None)
            self.assertEqual(
                mock.responses[0].getopt(pycurl.SSL_VERIFYHOST),
                None)

    def test_refresh_signature(self):
        blk_digest = '6f5902ac237024bdd0c176cb93063dc4+11'
        blk_sig = 'da39a3ee5e6b4b0d3255bfef95601890afd80709@53bed294'
        local_loc = blk_digest+'+A'+blk_sig
        remote_loc = blk_digest+'+R'+blk_sig
        api_client = self.mock_keep_services(count=1)
        headers = {'X-Keep-Locator':local_loc}
        with tutil.mock_keep_responses('', 200, **headers):
            # Check that the translated locator gets returned
            keep_client = arvados.KeepClient(api_client=api_client, block_cache=self.make_block_cache(self.disk_cache))
            self.assertEqual(local_loc, keep_client.refresh_signature(remote_loc))
            # Check that refresh_signature() uses the correct method and headers
            keep_client._get_or_head = mock.MagicMock()
            keep_client.refresh_signature(remote_loc)
            args, kwargs = keep_client._get_or_head.call_args_list[0]
            self.assertIn(remote_loc, args)
            self.assertEqual("HEAD", kwargs['method'])
            self.assertIn('X-Keep-Signature', kwargs['headers'])

    # test_*_timeout verify that KeepClient instructs pycurl to use
    # the appropriate connection and read timeouts. They don't care
    # whether pycurl actually exhibits the expected timeout behavior
    # -- those tests are in the KeepClientTimeout test class.

    def test_get_timeout(self):
        api_client = self.mock_keep_services(count=1)
        force_timeout = socket.timeout("timed out")
        with tutil.mock_keep_responses(force_timeout, 0) as mock:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            with self.assertRaises(arvados.errors.KeepReadError):
                keep_client.get('ffffffffffffffffffffffffffffffff')
            self.assertEqual(
                mock.responses[0].getopt(pycurl.CONNECTTIMEOUT_MS),
                int(arvados.KeepClient.DEFAULT_TIMEOUT[0]*1000))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_TIME),
                int(arvados.KeepClient.DEFAULT_TIMEOUT[1]))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_LIMIT),
                int(arvados.KeepClient.DEFAULT_TIMEOUT[2]))

    def test_put_timeout(self):
        api_client = self.mock_keep_services(count=1)
        force_timeout = socket.timeout("timed out")
        with tutil.mock_keep_responses(force_timeout, 0) as mock:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            with self.assertRaises(arvados.errors.KeepWriteError):
                keep_client.put(b'foo')
            self.assertEqual(
                mock.responses[0].getopt(pycurl.CONNECTTIMEOUT_MS),
                int(arvados.KeepClient.DEFAULT_TIMEOUT[0]*1000))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_TIME),
                int(arvados.KeepClient.DEFAULT_TIMEOUT[1]))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_LIMIT),
                int(arvados.KeepClient.DEFAULT_TIMEOUT[2]))

    def test_head_timeout(self):
        api_client = self.mock_keep_services(count=1)
        force_timeout = socket.timeout("timed out")
        with tutil.mock_keep_responses(force_timeout, 0) as mock:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            with self.assertRaises(arvados.errors.KeepReadError):
                keep_client.head('ffffffffffffffffffffffffffffffff')
            self.assertEqual(
                mock.responses[0].getopt(pycurl.CONNECTTIMEOUT_MS),
                int(arvados.KeepClient.DEFAULT_TIMEOUT[0]*1000))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_TIME),
                None)
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_LIMIT),
                None)

    def test_proxy_get_timeout(self):
        api_client = self.mock_keep_services(service_type='proxy', count=1)
        force_timeout = socket.timeout("timed out")
        with tutil.mock_keep_responses(force_timeout, 0) as mock:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            with self.assertRaises(arvados.errors.KeepReadError):
                keep_client.get('ffffffffffffffffffffffffffffffff')
            self.assertEqual(
                mock.responses[0].getopt(pycurl.CONNECTTIMEOUT_MS),
                int(arvados.KeepClient.DEFAULT_PROXY_TIMEOUT[0]*1000))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_TIME),
                int(arvados.KeepClient.DEFAULT_PROXY_TIMEOUT[1]))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_LIMIT),
                int(arvados.KeepClient.DEFAULT_PROXY_TIMEOUT[2]))

    def test_proxy_head_timeout(self):
        api_client = self.mock_keep_services(service_type='proxy', count=1)
        force_timeout = socket.timeout("timed out")
        with tutil.mock_keep_responses(force_timeout, 0) as mock:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            with self.assertRaises(arvados.errors.KeepReadError):
                keep_client.head('ffffffffffffffffffffffffffffffff')
            self.assertEqual(
                mock.responses[0].getopt(pycurl.CONNECTTIMEOUT_MS),
                int(arvados.KeepClient.DEFAULT_PROXY_TIMEOUT[0]*1000))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_TIME),
                None)
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_LIMIT),
                None)

    def test_proxy_put_timeout(self):
        self.disk_cache_dir = None
        api_client = self.mock_keep_services(service_type='proxy', count=1)
        force_timeout = socket.timeout("timed out")
        with tutil.mock_keep_responses(force_timeout, 0) as mock:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                num_retries=0,
            )
            with self.assertRaises(arvados.errors.KeepWriteError):
                keep_client.put('foo')
            self.assertEqual(
                mock.responses[0].getopt(pycurl.CONNECTTIMEOUT_MS),
                int(arvados.KeepClient.DEFAULT_PROXY_TIMEOUT[0]*1000))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_TIME),
                int(arvados.KeepClient.DEFAULT_PROXY_TIMEOUT[1]))
            self.assertEqual(
                mock.responses[0].getopt(pycurl.LOW_SPEED_LIMIT),
                int(arvados.KeepClient.DEFAULT_PROXY_TIMEOUT[2]))

    def check_no_services_error(self, verb, exc_class):
        api_client = mock.MagicMock(name='api_client')
        api_client.keep_services().accessible().execute.side_effect = (
            arvados.errors.ApiError)
        keep_client = arvados.KeepClient(
            api_client=api_client,
            block_cache=self.make_block_cache(self.disk_cache),
            num_retries=0,
        )
        with self.assertRaises(exc_class) as err_check:
            getattr(keep_client, verb)('d41d8cd98f00b204e9800998ecf8427e+0')
        self.assertEqual(0, len(err_check.exception.request_errors()))

    def test_get_error_with_no_services(self):
        self.check_no_services_error('get', arvados.errors.KeepReadError)

    def test_head_error_with_no_services(self):
        self.check_no_services_error('head', arvados.errors.KeepReadError)

    def test_put_error_with_no_services(self):
        self.check_no_services_error('put', arvados.errors.KeepWriteError)

    def check_errors_from_last_retry(self, verb, exc_class):
        api_client = self.mock_keep_services(count=2)
        req_mock = tutil.mock_keep_responses(
            "retry error reporting test", 500, 500, 500, 500, 500, 500, 502, 502)
        with req_mock, tutil.skip_sleep, \
                self.assertRaises(exc_class) as err_check:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            getattr(keep_client, verb)('d41d8cd98f00b204e9800998ecf8427e+0',
                                       num_retries=3)
        self.assertEqual([502, 502], [
                getattr(error, 'status_code', None)
                for error in err_check.exception.request_errors().values()])
        self.assertRegex(str(err_check.exception), r'failed to (read|write) .* after 4 attempts')

    def test_get_error_reflects_last_retry(self):
        self.check_errors_from_last_retry('get', arvados.errors.KeepReadError)

    def test_head_error_reflects_last_retry(self):
        self.check_errors_from_last_retry('head', arvados.errors.KeepReadError)

    def test_put_error_reflects_last_retry(self):
        self.check_errors_from_last_retry('put', arvados.errors.KeepWriteError)

    def test_put_error_does_not_include_successful_puts(self):
        data = 'partial failure test'
        data_loc = tutil.str_keep_locator(data)
        api_client = self.mock_keep_services(count=3)
        with tutil.mock_keep_responses(data_loc, 200, 500, 500) as req_mock, \
                self.assertRaises(arvados.errors.KeepWriteError) as exc_check:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            keep_client.put(data)
        self.assertEqual(2, len(exc_check.exception.request_errors()))

    def test_proxy_put_with_no_writable_services(self):
        data = 'test with no writable services'
        data_loc = tutil.str_keep_locator(data)
        api_client = self.mock_keep_services(service_type='proxy', read_only=True, count=1)
        with tutil.mock_keep_responses(data_loc, 200, 500, 500) as req_mock, \
                self.assertRaises(arvados.errors.KeepWriteError) as exc_check:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            keep_client.put(data)
        self.assertEqual(True, ("no Keep services available" in str(exc_check.exception)))
        self.assertEqual(0, len(exc_check.exception.request_errors()))

    def test_oddball_service_get(self):
        body = b'oddball service get'
        api_client = self.mock_keep_services(service_type='fancynewblobstore')
        with tutil.mock_keep_responses(body, 200):
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            actual = keep_client.get(tutil.str_keep_locator(body))
        self.assertEqual(body, actual)

    def test_oddball_service_put(self):
        body = b'oddball service put'
        pdh = tutil.str_keep_locator(body)
        api_client = self.mock_keep_services(service_type='fancynewblobstore')
        with tutil.mock_keep_responses(pdh, 200):
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            actual = keep_client.put(body, copies=1)
        self.assertEqual(pdh, actual)

    def test_oddball_service_writer_count(self):
        body = b'oddball service writer count'
        pdh = tutil.str_keep_locator(body)
        api_client = self.mock_keep_services(service_type='fancynewblobstore',
                                             count=4)
        headers = {'x-keep-replicas-stored': 3}
        with tutil.mock_keep_responses(pdh, 200, 418, 418, 418,
                                       **headers) as req_mock:
            keep_client = arvados.KeepClient(
                api_client=api_client,
                block_cache=self.make_block_cache(self.disk_cache),
                num_retries=0,
            )
            actual = keep_client.put(body, copies=2)
        self.assertEqual(pdh, actual)
        self.assertEqual(1, req_mock.call_count)


@tutil.skip_sleep
@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepClientCacheTestCase(unittest.TestCase, tutil.ApiClientMock, DiskCacheBase):
    disk_cache = False

    def setUp(self):
        self.api_client = self.mock_keep_services(count=2)
        self.keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=self.make_block_cache(self.disk_cache))
        self.data = b'xyzzy'
        self.locator = '1271ed5ef305aadabc605b1609e24c52'

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    @mock.patch('arvados.KeepClient.KeepService.get')
    def test_get_request_cache(self, get_mock):
        with tutil.mock_keep_responses(self.data, 200, 200):
            self.keep_client.get(self.locator)
            self.keep_client.get(self.locator)
        # Request already cached, don't require more than one request
        get_mock.assert_called_once()

    @mock.patch('arvados.KeepClient.KeepService.get')
    def test_head_request_cache(self, get_mock):
        with tutil.mock_keep_responses(self.data, 200, 200):
            self.keep_client.head(self.locator)
            self.keep_client.head(self.locator)
        # Don't cache HEAD requests so that they're not confused with GET reqs
        self.assertEqual(2, get_mock.call_count)

    @mock.patch('arvados.KeepClient.KeepService.get')
    def test_head_and_then_get_return_different_responses(self, get_mock):
        head_resp = None
        get_resp = None
        get_mock.side_effect = [b'first response', b'second response']
        with tutil.mock_keep_responses(self.data, 200, 200):
            head_resp = self.keep_client.head(self.locator)
            get_resp = self.keep_client.get(self.locator)
        self.assertEqual(b'first response', head_resp)
        # First reponse was not cached because it was from a HEAD request.
        self.assertNotEqual(head_resp, get_resp)


@tutil.skip_sleep
@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepXRequestIdTestCase(unittest.TestCase, tutil.ApiClientMock, DiskCacheBase):
    disk_cache = False

    def setUp(self):
        self.api_client = self.mock_keep_services(count=2)
        self.keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=self.make_block_cache(self.disk_cache))
        self.data = b'xyzzy'
        self.locator = '1271ed5ef305aadabc605b1609e24c52'
        self.test_id = arvados.util.new_request_id()
        self.assertRegex(self.test_id, r'^req-[a-z0-9]{20}$')
        # If we don't set request_id to None explicitly here, it will
        # return <MagicMock name='api_client_mock.request_id'
        # id='123456789'>:
        self.api_client.request_id = None

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def test_default_to_api_client_request_id(self):
        self.api_client.request_id = self.test_id
        with tutil.mock_keep_responses(self.locator, 200, 200) as mock:
            self.keep_client.put(self.data)
        self.assertEqual(2, len(mock.responses))
        for resp in mock.responses:
            self.assertProvidedRequestId(resp)

        with tutil.mock_keep_responses(self.data, 200) as mock:
            self.keep_client.get(self.locator)
        self.assertProvidedRequestId(mock.responses[0])

        with tutil.mock_keep_responses(b'', 200) as mock:
            self.keep_client.head(self.locator)
        self.assertProvidedRequestId(mock.responses[0])

    def test_explicit_request_id(self):
        with tutil.mock_keep_responses(self.locator, 200, 200) as mock:
            self.keep_client.put(self.data, request_id=self.test_id)
        self.assertEqual(2, len(mock.responses))
        for resp in mock.responses:
            self.assertProvidedRequestId(resp)

        with tutil.mock_keep_responses(self.data, 200) as mock:
            self.keep_client.get(self.locator, request_id=self.test_id)
        self.assertProvidedRequestId(mock.responses[0])

        with tutil.mock_keep_responses(b'', 200) as mock:
            self.keep_client.head(self.locator, request_id=self.test_id)
        self.assertProvidedRequestId(mock.responses[0])

    def test_automatic_request_id(self):
        with tutil.mock_keep_responses(self.locator, 200, 200) as mock:
            self.keep_client.put(self.data)
        self.assertEqual(2, len(mock.responses))
        for resp in mock.responses:
            self.assertAutomaticRequestId(resp)

        with tutil.mock_keep_responses(self.data, 200) as mock:
            self.keep_client.get(self.locator)
        self.assertAutomaticRequestId(mock.responses[0])

        with tutil.mock_keep_responses(b'', 200) as mock:
            self.keep_client.head(self.locator)
        self.assertAutomaticRequestId(mock.responses[0])

    def test_request_id_in_exception(self):
        with tutil.mock_keep_responses(b'', 400, 400, 400) as mock:
            with self.assertRaisesRegex(arvados.errors.KeepReadError, self.test_id):
                self.keep_client.head(self.locator, request_id=self.test_id)

        with tutil.mock_keep_responses(b'', 400, 400, 400) as mock:
            with self.assertRaisesRegex(arvados.errors.KeepReadError, r'req-[a-z0-9]{20}'):
                self.keep_client.get(self.locator)

        with tutil.mock_keep_responses(b'', 400, 400, 400) as mock:
            with self.assertRaisesRegex(arvados.errors.KeepWriteError, self.test_id):
                self.keep_client.put(self.data, request_id=self.test_id)

        with tutil.mock_keep_responses(b'', 400, 400, 400) as mock:
            with self.assertRaisesRegex(arvados.errors.KeepWriteError, r'req-[a-z0-9]{20}'):
                self.keep_client.put(self.data)

    def assertAutomaticRequestId(self, resp):
        hdr = [x for x in resp.getopt(pycurl.HTTPHEADER)
               if x.startswith('X-Request-Id: ')][0]
        self.assertNotEqual(hdr, 'X-Request-Id: '+self.test_id)
        self.assertRegex(hdr, r'^X-Request-Id: req-[a-z0-9]{20}$')

    def assertProvidedRequestId(self, resp):
        self.assertIn('X-Request-Id: '+self.test_id,
                      resp.getopt(pycurl.HTTPHEADER))


@tutil.skip_sleep
#@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepClientRendezvousTestCase(unittest.TestCase, tutil.ApiClientMock, DiskCacheBase):
    disk_cache = False

    def setUp(self):
        # expected_order[i] is the probe order for
        # hash=md5(sprintf("%064x",i)) where there are 16 services
        # with uuid sprintf("anything-%015x",j) with j in 0..15. E.g.,
        # the first probe for the block consisting of 64 "0"
        # characters is the service whose uuid is
        # "zzzzz-bi6l4-000000000000003", so expected_order[0][0]=='3'.
        self.services = 16
        self.expected_order = [
            list('3eab2d5fc9681074'),
            list('097dba52e648f1c3'),
            list('c5b4e023f8a7d691'),
            list('9d81c02e76a3bf54'),
            ]
        self.blocks = [
            "{:064x}".format(x).encode()
            for x in range(len(self.expected_order))]
        self.hashes = [
            hashlib.md5(self.blocks[x]).hexdigest()
            for x in range(len(self.expected_order))]
        self.api_client = self.mock_keep_services(count=self.services)
        self.keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=self.make_block_cache(self.disk_cache))

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def test_weighted_service_roots_against_reference_set(self):
        # Confirm weighted_service_roots() returns the correct order
        for i, hash in enumerate(self.hashes):
            roots = self.keep_client.weighted_service_roots(arvados.KeepLocator(hash))
            got_order = [
                re.search(r'//\[?keep0x([0-9a-f]+)', root).group(1)
                for root in roots]
            self.assertEqual(self.expected_order[i], got_order)

    def test_get_probe_order_against_reference_set(self):
        self._test_probe_order_against_reference_set(
            lambda i: self.keep_client.get(self.hashes[i], num_retries=1))

    def test_head_probe_order_against_reference_set(self):
        self._test_probe_order_against_reference_set(
            lambda i: self.keep_client.head(self.hashes[i], num_retries=1))

    def test_put_probe_order_against_reference_set(self):
        # copies=1 prevents the test from being sensitive to races
        # between writer threads.
        self._test_probe_order_against_reference_set(
            lambda i: self.keep_client.put(self.blocks[i], num_retries=1, copies=1))

    def _test_probe_order_against_reference_set(self, op):
        for i in range(len(self.blocks)):
            with tutil.mock_keep_responses('', *[500 for _ in range(self.services*2)]) as mock, \
                 self.assertRaises(arvados.errors.KeepRequestError):
                op(i)
            got_order = [
                re.search(r'//\[?keep0x([0-9a-f]+)', resp.getopt(pycurl.URL).decode()).group(1)
                for resp in mock.responses]
            self.assertEqual(self.expected_order[i]*2, got_order)

    def test_put_probe_order_multiple_copies(self):
        for copies in range(2, 4):
            for i in range(len(self.blocks)):
                with tutil.mock_keep_responses('', *[500 for _ in range(self.services*3)]) as mock, \
                     self.assertRaises(arvados.errors.KeepWriteError):
                    self.keep_client.put(self.blocks[i], num_retries=2, copies=copies)
                got_order = [
                    re.search(r'//\[?keep0x([0-9a-f]+)', resp.getopt(pycurl.URL).decode()).group(1)
                    for resp in mock.responses]
                # With T threads racing to make requests, the position
                # of a given server in the sequence of HTTP requests
                # (got_order) cannot be more than T-1 positions
                # earlier than that server's position in the reference
                # probe sequence (expected_order).
                #
                # Loop invariant: we have accounted for +pos+ expected
                # probes, either by seeing them in +got_order+ or by
                # putting them in +pending+ in the hope of seeing them
                # later. As long as +len(pending)<T+, we haven't
                # started a request too early.
                pending = []
                for pos, expected in enumerate(self.expected_order[i]*3):
                    got = got_order[pos-len(pending)]
                    while got in pending:
                        del pending[pending.index(got)]
                        got = got_order[pos-len(pending)]
                    if got != expected:
                        pending.append(expected)
                        self.assertLess(
                            len(pending), copies,
                            "pending={}, with copies={}, got {}, expected {}".format(
                                pending, copies, repr(got_order), repr(self.expected_order[i]*3)))

    def test_probe_waste_adding_one_server(self):
        hashes = [
            hashlib.md5("{:064x}".format(x).encode()).hexdigest() for x in range(100)]
        initial_services = 12
        self.api_client = self.mock_keep_services(count=initial_services)
        self.keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=self.make_block_cache(self.disk_cache))
        probes_before = [
            self.keep_client.weighted_service_roots(arvados.KeepLocator(hash)) for hash in hashes]
        for added_services in range(1, 12):
            api_client = self.mock_keep_services(count=initial_services+added_services)
            keep_client = arvados.KeepClient(api_client=api_client, block_cache=self.make_block_cache(self.disk_cache))
            total_penalty = 0
            for hash_index in range(len(hashes)):
                probe_after = keep_client.weighted_service_roots(
                    arvados.KeepLocator(hashes[hash_index]))
                penalty = probe_after.index(probes_before[hash_index][0])
                self.assertLessEqual(penalty, added_services)
                total_penalty += penalty
            # Average penalty per block should not exceed
            # N(added)/N(orig) by more than 20%, and should get closer
            # to the ideal as we add data points.
            expect_penalty = (
                added_services *
                len(hashes) / initial_services)
            max_penalty = (
                expect_penalty *
                (120 - added_services)/100)
            min_penalty = (
                expect_penalty * 8/10)
            self.assertTrue(
                min_penalty <= total_penalty <= max_penalty,
                "With {}+{} services, {} blocks, penalty {} but expected {}..{}".format(
                    initial_services,
                    added_services,
                    len(hashes),
                    total_penalty,
                    min_penalty,
                    max_penalty))

    def check_64_zeros_error_order(self, verb, exc_class):
        data = b'0' * 64
        if verb == 'get':
            data = tutil.str_keep_locator(data)
        # Arbitrary port number:
        aport = random.randint(1024,65535)
        api_client = self.mock_keep_services(service_port=aport, count=self.services)
        keep_client = arvados.KeepClient(api_client=api_client, block_cache=self.make_block_cache(self.disk_cache))
        with mock.patch('pycurl.Curl') as curl_mock, \
             self.assertRaises(exc_class) as err_check:
            curl_mock.return_value = tutil.FakeCurl.make(code=500, body=b'')
            getattr(keep_client, verb)(data)
        urls = [urllib.parse.urlparse(url)
                for url in err_check.exception.request_errors()]
        self.assertEqual([('keep0x' + c, aport) for c in '3eab2d5fc9681074'],
                         [(url.hostname, url.port) for url in urls])

    def test_get_error_shows_probe_order(self):
        self.check_64_zeros_error_order('get', arvados.errors.KeepReadError)

    def test_put_error_shows_probe_order(self):
        self.check_64_zeros_error_order('put', arvados.errors.KeepWriteError)


@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepClientTimeout(keepstub.StubKeepServers, unittest.TestCase, DiskCacheBase):
    disk_cache = False

    # BANDWIDTH_LOW_LIM must be less than len(DATA) so we can transfer
    # 1s worth of data and then trigger bandwidth errors before running
    # out of data.
    DATA = b'x'*2**11
    BANDWIDTH_LOW_LIM = 1024
    TIMEOUT_TIME = 1.0

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    class assertTakesBetween(unittest.TestCase):
        def __init__(self, tmin, tmax):
            self.tmin = tmin
            self.tmax = tmax

        def __enter__(self):
            self.t0 = time.time()

        def __exit__(self, *args, **kwargs):
            # Round times to milliseconds, like CURL. Otherwise, we
            # fail when CURL reaches a 1s timeout at 0.9998s.
            delta = round(time.time() - self.t0, 3)
            self.assertGreaterEqual(delta, self.tmin)
            self.assertLessEqual(delta, self.tmax)

    class assertTakesGreater(unittest.TestCase):
        def __init__(self, tmin):
            self.tmin = tmin

        def __enter__(self):
            self.t0 = time.time()

        def __exit__(self, *args, **kwargs):
            delta = round(time.time() - self.t0, 3)
            self.assertGreaterEqual(delta, self.tmin)

    def keepClient(self, timeouts=(0.1, TIMEOUT_TIME, BANDWIDTH_LOW_LIM)):
        return arvados.KeepClient(
            api_client=self.api_client,
            timeout=timeouts, block_cache=self.make_block_cache(self.disk_cache))

    def test_timeout_slow_connect(self):
        # Can't simulate TCP delays with our own socket. Leave our
        # stub server running uselessly, and try to connect to an
        # unroutable IP address instead.
        self.api_client = self.mock_keep_services(
            count=1,
            service_host='240.0.0.0',
        )
        with self.assertTakesBetween(0.1, 0.5):
            with self.assertRaises(arvados.errors.KeepWriteError):
                self.keepClient().put(self.DATA, copies=1, num_retries=0)

    def test_low_bandwidth_no_delays_success(self):
        self.server.setbandwidth(2*self.BANDWIDTH_LOW_LIM)
        kc = self.keepClient()
        loc = kc.put(self.DATA, copies=1, num_retries=0)
        self.assertEqual(self.DATA, kc.get(loc, num_retries=0))

    def test_too_low_bandwidth_no_delays_failure(self):
        # Check that lessening bandwidth corresponds to failing
        kc = self.keepClient()
        loc = kc.put(self.DATA, copies=1, num_retries=0)
        self.server.setbandwidth(0.5*self.BANDWIDTH_LOW_LIM)
        with self.assertTakesGreater(self.TIMEOUT_TIME):
            with self.assertRaises(arvados.errors.KeepReadError):
                kc.get(loc, num_retries=0)
        with self.assertTakesGreater(self.TIMEOUT_TIME):
            with self.assertRaises(arvados.errors.KeepWriteError):
                kc.put(self.DATA, copies=1, num_retries=0)

    def test_low_bandwidth_with_server_response_delay_failure(self):
        kc = self.keepClient()
        loc = kc.put(self.DATA, copies=1, num_retries=0)
        self.server.setbandwidth(self.BANDWIDTH_LOW_LIM)
        # Note the actual delay must be 1s longer than the low speed
        # limit interval in order for curl to detect it reliably.
        self.server.setdelays(response=self.TIMEOUT_TIME+1)
        with self.assertTakesGreater(self.TIMEOUT_TIME):
            with self.assertRaises(arvados.errors.KeepReadError):
                kc.get(loc, num_retries=0)
        with self.assertTakesGreater(self.TIMEOUT_TIME):
            with self.assertRaises(arvados.errors.KeepWriteError):
                kc.put(self.DATA, copies=1, num_retries=0)
        with self.assertTakesGreater(self.TIMEOUT_TIME):
            kc.head(loc, num_retries=0)

    def test_low_bandwidth_with_server_mid_delay_failure(self):
        kc = self.keepClient()
        loc = kc.put(self.DATA, copies=1, num_retries=0)
        self.server.setbandwidth(self.BANDWIDTH_LOW_LIM)
        # Note the actual delay must be 1s longer than the low speed
        # limit interval in order for curl to detect it reliably.
        self.server.setdelays(mid_write=self.TIMEOUT_TIME+1, mid_read=self.TIMEOUT_TIME+1)
        with self.assertTakesGreater(self.TIMEOUT_TIME):
            with self.assertRaises(arvados.errors.KeepReadError) as e:
                kc.get(loc, num_retries=0)
        with self.assertTakesGreater(self.TIMEOUT_TIME):
            with self.assertRaises(arvados.errors.KeepWriteError):
                kc.put(self.DATA, copies=1, num_retries=0)

    def test_timeout_slow_request(self):
        loc = self.keepClient().put(self.DATA, copies=1, num_retries=0)
        self.server.setdelays(request=.2)
        self._test_connect_timeout_under_200ms(loc)
        self.server.setdelays(request=2)
        self._test_response_timeout_under_2s(loc)

    def test_timeout_slow_response(self):
        loc = self.keepClient().put(self.DATA, copies=1, num_retries=0)
        self.server.setdelays(response=.2)
        self._test_connect_timeout_under_200ms(loc)
        self.server.setdelays(response=2)
        self._test_response_timeout_under_2s(loc)

    def test_timeout_slow_response_body(self):
        loc = self.keepClient().put(self.DATA, copies=1, num_retries=0)
        self.server.setdelays(response_body=.2)
        self._test_connect_timeout_under_200ms(loc)
        self.server.setdelays(response_body=2)
        self._test_response_timeout_under_2s(loc)

    def _test_connect_timeout_under_200ms(self, loc):
        # Allow 100ms to connect, then 1s for response. Everything
        # should work, and everything should take at least 200ms to
        # return.
        kc = self.keepClient(timeouts=(.1, 1))
        with self.assertTakesBetween(.2, .3):
            kc.put(self.DATA, copies=1, num_retries=0)
        with self.assertTakesBetween(.2, .3):
            self.assertEqual(self.DATA, kc.get(loc, num_retries=0))

    def _test_response_timeout_under_2s(self, loc):
        # Allow 10s to connect, then 1s for response. Nothing should
        # work, and everything should take at least 1s to return.
        kc = self.keepClient(timeouts=(10, 1))
        with self.assertTakesBetween(1, 9):
            with self.assertRaises(arvados.errors.KeepReadError):
                kc.get(loc, num_retries=0)
        with self.assertTakesBetween(1, 9):
            with self.assertRaises(arvados.errors.KeepWriteError):
                kc.put(self.DATA, copies=1, num_retries=0)


@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepClientGatewayTestCase(unittest.TestCase, tutil.ApiClientMock, DiskCacheBase):
    disk_cache = False

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def mock_disks_and_gateways(self, disks=3, gateways=1):
        self.gateways = [{
                'uuid': 'zzzzz-bi6l4-gateway{:08d}'.format(i),
                'owner_uuid': 'zzzzz-tpzed-000000000000000',
                'service_host': 'gatewayhost{}'.format(i),
                'service_port': 12345,
                'service_ssl_flag': True,
                'service_type': 'gateway:test',
        } for i in range(gateways)]
        self.gateway_roots = [
            "https://{service_host}:{service_port}/".format(**gw)
            for gw in self.gateways]
        self.api_client = self.mock_keep_services(
            count=disks, additional_services=self.gateways)
        self.keepClient = arvados.KeepClient(api_client=self.api_client, block_cache=self.make_block_cache(self.disk_cache))

    @mock.patch('pycurl.Curl')
    def test_get_with_gateway_hint_first(self, MockCurl):
        MockCurl.return_value = tutil.FakeCurl.make(
            code=200, body='foo', headers={'Content-Length': 3})
        self.mock_disks_and_gateways()
        locator = 'acbd18db4cc2f85cedef654fccc4a4d8+3+K@' + self.gateways[0]['uuid']
        self.assertEqual(b'foo', self.keepClient.get(locator))
        self.assertEqual(self.gateway_roots[0]+locator,
                         MockCurl.return_value.getopt(pycurl.URL).decode())
        self.assertEqual(True, self.keepClient.head(locator))

    @mock.patch('pycurl.Curl')
    def test_get_with_gateway_hints_in_order(self, MockCurl):
        gateways = 4
        disks = 3
        mocks = [
            tutil.FakeCurl.make(code=404, body='')
            for _ in range(gateways+disks)
        ]
        MockCurl.side_effect = tutil.queue_with(mocks)
        self.mock_disks_and_gateways(gateways=gateways, disks=disks)
        locator = '+'.join(['acbd18db4cc2f85cedef654fccc4a4d8+3'] +
                           ['K@'+gw['uuid'] for gw in self.gateways])
        with self.assertRaises(arvados.errors.NotFoundError):
            self.keepClient.get(locator)
        # Gateways are tried first, in the order given.
        for i, root in enumerate(self.gateway_roots):
            self.assertEqual(root+locator,
                             mocks[i].getopt(pycurl.URL).decode())
        # Disk services are tried next.
        for i in range(gateways, gateways+disks):
            self.assertRegex(
                mocks[i].getopt(pycurl.URL).decode(),
                r'keep0x')

    @mock.patch('pycurl.Curl')
    def test_head_with_gateway_hints_in_order(self, MockCurl):
        gateways = 4
        disks = 3
        mocks = [
            tutil.FakeCurl.make(code=404, body=b'')
            for _ in range(gateways+disks)
        ]
        MockCurl.side_effect = tutil.queue_with(mocks)
        self.mock_disks_and_gateways(gateways=gateways, disks=disks)
        locator = '+'.join(['acbd18db4cc2f85cedef654fccc4a4d8+3'] +
                           ['K@'+gw['uuid'] for gw in self.gateways])
        with self.assertRaises(arvados.errors.NotFoundError):
            self.keepClient.head(locator)
        # Gateways are tried first, in the order given.
        for i, root in enumerate(self.gateway_roots):
            self.assertEqual(root+locator,
                             mocks[i].getopt(pycurl.URL).decode())
        # Disk services are tried next.
        for i in range(gateways, gateways+disks):
            self.assertRegex(
                mocks[i].getopt(pycurl.URL).decode(),
                r'keep0x')

    @mock.patch('pycurl.Curl')
    def test_get_with_remote_proxy_hint(self, MockCurl):
        MockCurl.return_value = tutil.FakeCurl.make(
            code=200, body=b'foo', headers={'Content-Length': 3})
        self.mock_disks_and_gateways()
        locator = 'acbd18db4cc2f85cedef654fccc4a4d8+3+K@xyzzy'
        self.assertEqual(b'foo', self.keepClient.get(locator))
        self.assertEqual('https://keep.xyzzy.arvadosapi.com/'+locator,
                         MockCurl.return_value.getopt(pycurl.URL).decode())

    @mock.patch('pycurl.Curl')
    def test_head_with_remote_proxy_hint(self, MockCurl):
        MockCurl.return_value = tutil.FakeCurl.make(
            code=200, body=b'foo', headers={'Content-Length': 3})
        self.mock_disks_and_gateways()
        locator = 'acbd18db4cc2f85cedef654fccc4a4d8+3+K@xyzzy'
        self.assertEqual(True, self.keepClient.head(locator))
        self.assertEqual('https://keep.xyzzy.arvadosapi.com/'+locator,
                         MockCurl.return_value.getopt(pycurl.URL).decode())


class KeepClientRetryTestMixin(object):
    disk_cache = False

    # Testing with a local Keep store won't exercise the retry behavior.
    # Instead, our strategy is:
    # * Create a client with one proxy specified (pointed at a black
    #   hole), so there's no need to instantiate an API client, and
    #   all HTTP requests come from one place.
    # * Mock httplib's request method to provide simulated responses.
    # This lets us test the retry logic extensively without relying on any
    # supporting servers, and prevents side effects in case something hiccups.
    # To use this mixin, define DEFAULT_EXPECT, DEFAULT_EXCEPTION, and
    # run_method().
    #
    # Test classes must define TEST_PATCHER to a method that mocks
    # out appropriate methods in the client.

    PROXY_ADDR = 'http://[%s]:65535/' % (tutil.TEST_HOST,)
    TEST_DATA = b'testdata'
    TEST_LOCATOR = 'ef654c40ab4f1747fc699915d4f70902+8'

    def setUp(self):
        self.client_kwargs = {'proxy': self.PROXY_ADDR, 'local_store': ''}

    def new_client(self, **caller_kwargs):
        kwargs = self.client_kwargs.copy()
        kwargs.update(caller_kwargs)
        kwargs['block_cache'] = self.make_block_cache(self.disk_cache)
        return arvados.KeepClient(**kwargs)

    def run_method(self, *args, **kwargs):
        raise NotImplementedError("test subclasses must define run_method")

    def check_success(self, expected=None, *args, **kwargs):
        if expected is None:
            expected = self.DEFAULT_EXPECT
        self.assertEqual(expected, self.run_method(*args, **kwargs))

    def check_exception(self, error_class=None, *args, **kwargs):
        if error_class is None:
            error_class = self.DEFAULT_EXCEPTION
        with self.assertRaises(error_class) as err:
            self.run_method(*args, **kwargs)
        return err

    def test_immediate_success(self):
        with self.TEST_PATCHER(self.DEFAULT_EXPECT, 200):
            self.check_success()

    def test_retry_then_success(self):
        with self.TEST_PATCHER(self.DEFAULT_EXPECT, 500, 200):
            self.check_success(num_retries=3)

    def test_exception_then_success(self):
        with self.TEST_PATCHER(self.DEFAULT_EXPECT, Exception('mock err'), 200):
            self.check_success(num_retries=3)

    def test_no_retry_after_permanent_error(self):
        with self.TEST_PATCHER(self.DEFAULT_EXPECT, 403, 200):
            self.check_exception(num_retries=3)

    def test_error_after_retries_exhausted(self):
        with self.TEST_PATCHER(self.DEFAULT_EXPECT, 500, 500, 200):
            err = self.check_exception(num_retries=1)
        self.assertRegex(str(err.exception), r'failed to .* after 2 attempts')

    def test_num_retries_instance_fallback(self):
        self.client_kwargs['num_retries'] = 3
        with self.TEST_PATCHER(self.DEFAULT_EXPECT, 500, 200):
            self.check_success()


@tutil.skip_sleep
@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepClientRetryGetTestCase(KeepClientRetryTestMixin, unittest.TestCase, DiskCacheBase):
    DEFAULT_EXPECT = KeepClientRetryTestMixin.TEST_DATA
    DEFAULT_EXCEPTION = arvados.errors.KeepReadError
    HINTED_LOCATOR = KeepClientRetryTestMixin.TEST_LOCATOR + '+K@xyzzy'
    TEST_PATCHER = staticmethod(tutil.mock_keep_responses)

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def run_method(self, locator=KeepClientRetryTestMixin.TEST_LOCATOR,
                   *args, **kwargs):
        return self.new_client().get(locator, *args, **kwargs)

    def test_specific_exception_when_not_found(self):
        with tutil.mock_keep_responses(self.DEFAULT_EXPECT, 404, 200):
            self.check_exception(arvados.errors.NotFoundError, num_retries=3)

    def test_general_exception_with_mixed_errors(self):
        # get should raise a NotFoundError if no server returns the block,
        # and a high threshold of servers report that it's not found.
        # This test rigs up 50/50 disagreement between two servers, and
        # checks that it does not become a NotFoundError.
        client = self.new_client(num_retries=0)
        with tutil.mock_keep_responses(self.DEFAULT_EXPECT, 404, 500):
            with self.assertRaises(arvados.errors.KeepReadError) as exc_check:
                client.get(self.HINTED_LOCATOR)
            self.assertNotIsInstance(
                exc_check.exception, arvados.errors.NotFoundError,
                "mixed errors raised NotFoundError")

    def test_hint_server_can_succeed_without_retries(self):
        with tutil.mock_keep_responses(self.DEFAULT_EXPECT, 404, 200, 500):
            self.check_success(locator=self.HINTED_LOCATOR)

    def test_try_next_server_after_timeout(self):
        with tutil.mock_keep_responses(
                (socket.timeout("timed out"), 200),
                (self.DEFAULT_EXPECT, 200)):
            self.check_success(locator=self.HINTED_LOCATOR)

    def test_retry_data_with_wrong_checksum(self):
        with tutil.mock_keep_responses(
                ('baddata', 200),
                (self.DEFAULT_EXPECT, 200)):
            self.check_success(locator=self.HINTED_LOCATOR)


@tutil.skip_sleep
@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepClientRetryHeadTestCase(KeepClientRetryTestMixin, unittest.TestCase, DiskCacheBase):
    DEFAULT_EXPECT = True
    DEFAULT_EXCEPTION = arvados.errors.KeepReadError
    HINTED_LOCATOR = KeepClientRetryTestMixin.TEST_LOCATOR + '+K@xyzzy'
    TEST_PATCHER = staticmethod(tutil.mock_keep_responses)

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def run_method(self, locator=KeepClientRetryTestMixin.TEST_LOCATOR,
                   *args, **kwargs):
        return self.new_client().head(locator, *args, **kwargs)

    def test_specific_exception_when_not_found(self):
        with tutil.mock_keep_responses(self.DEFAULT_EXPECT, 404, 200):
            self.check_exception(arvados.errors.NotFoundError, num_retries=3)

    def test_general_exception_with_mixed_errors(self):
        # head should raise a NotFoundError if no server returns the block,
        # and a high threshold of servers report that it's not found.
        # This test rigs up 50/50 disagreement between two servers, and
        # checks that it does not become a NotFoundError.
        client = self.new_client(num_retries=0)
        with tutil.mock_keep_responses(self.DEFAULT_EXPECT, 404, 500):
            with self.assertRaises(arvados.errors.KeepReadError) as exc_check:
                client.head(self.HINTED_LOCATOR)
            self.assertNotIsInstance(
                exc_check.exception, arvados.errors.NotFoundError,
                "mixed errors raised NotFoundError")

    def test_hint_server_can_succeed_without_retries(self):
        with tutil.mock_keep_responses(self.DEFAULT_EXPECT, 404, 200, 500):
            self.check_success(locator=self.HINTED_LOCATOR)

    def test_try_next_server_after_timeout(self):
        with tutil.mock_keep_responses(
                (socket.timeout("timed out"), 200),
                (self.DEFAULT_EXPECT, 200)):
            self.check_success(locator=self.HINTED_LOCATOR)


@tutil.skip_sleep
@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepClientRetryPutTestCase(KeepClientRetryTestMixin, unittest.TestCase, DiskCacheBase):
    DEFAULT_EXPECT = KeepClientRetryTestMixin.TEST_LOCATOR
    DEFAULT_EXCEPTION = arvados.errors.KeepWriteError
    TEST_PATCHER = staticmethod(tutil.mock_keep_responses)

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def run_method(self, data=KeepClientRetryTestMixin.TEST_DATA,
                   copies=1, *args, **kwargs):
        return self.new_client().put(data, copies, *args, **kwargs)

    def test_do_not_send_multiple_copies_to_same_server(self):
        with tutil.mock_keep_responses(self.DEFAULT_EXPECT, 200):
            self.check_exception(copies=2, num_retries=3)


class AvoidOverreplication(unittest.TestCase, tutil.ApiClientMock):
    class FakeKeepService(object):
        def __init__(self, delay, will_succeed=False, will_raise=None, replicas=1):
            self.delay = delay
            self.will_succeed = will_succeed
            self.will_raise = will_raise
            self._result = {}
            self._result['headers'] = {}
            self._result['headers']['x-keep-replicas-stored'] = str(replicas)
            self._result['headers']['x-keep-storage-classes-confirmed'] = 'default={}'.format(replicas)
            self._result['body'] = 'foobar'

        def put(self, data_hash, data, timeout, headers):
            time.sleep(self.delay)
            if self.will_raise is not None:
                raise self.will_raise
            return self.will_succeed

        def last_result(self):
            if self.will_succeed:
                return self._result
            else:
                return {"status_code": 500, "body": "didn't succeed"}

        def finished(self):
            return False


    def setUp(self):
        self.copies = 3
        self.pool = arvados.KeepClient.KeepWriterThreadPool(
            data = 'foo',
            data_hash = 'acbd18db4cc2f85cedef654fccc4a4d8+3',
            max_service_replicas = self.copies,
            copies = self.copies
        )

    def test_only_write_enough_on_success(self):
        for i in range(10):
            ks = self.FakeKeepService(delay=i/10.0, will_succeed=True)
            self.pool.add_task(ks, None)
        self.pool.join()
        self.assertEqual(self.pool.done(), (self.copies, []))

    def test_only_write_enough_on_partial_success(self):
        for i in range(5):
            ks = self.FakeKeepService(delay=i/10.0, will_succeed=False)
            self.pool.add_task(ks, None)
            ks = self.FakeKeepService(delay=i/10.0, will_succeed=True)
            self.pool.add_task(ks, None)
        self.pool.join()
        self.assertEqual(self.pool.done(), (self.copies, []))

    def test_only_write_enough_when_some_crash(self):
        for i in range(5):
            ks = self.FakeKeepService(delay=i/10.0, will_raise=Exception())
            self.pool.add_task(ks, None)
            ks = self.FakeKeepService(delay=i/10.0, will_succeed=True)
            self.pool.add_task(ks, None)
        self.pool.join()
        self.assertEqual(self.pool.done(), (self.copies, []))

    def test_fail_when_too_many_crash(self):
        for i in range(self.copies+1):
            ks = self.FakeKeepService(delay=i/10.0, will_raise=Exception())
            self.pool.add_task(ks, None)
        for i in range(self.copies-1):
            ks = self.FakeKeepService(delay=i/10.0, will_succeed=True)
            self.pool.add_task(ks, None)
        self.pool.join()
        self.assertEqual(self.pool.done(), (self.copies-1, []))


@tutil.skip_sleep
@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class RetryNeedsMultipleServices(unittest.TestCase, tutil.ApiClientMock, DiskCacheBase):
    block_cache = False

    # Test put()s that need two distinct servers to succeed, possibly
    # requiring multiple passes through the retry loop.

    def setUp(self):
        self.api_client = self.mock_keep_services(count=2)
        self.keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=self.make_block_cache(self.disk_cache))

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def test_success_after_exception(self):
        with tutil.mock_keep_responses(
                'acbd18db4cc2f85cedef654fccc4a4d8+3',
                Exception('mock err'), 200, 200) as req_mock:
            self.keep_client.put('foo', num_retries=1, copies=2)
        self.assertEqual(3, req_mock.call_count)

    def test_success_after_retryable_error(self):
        with tutil.mock_keep_responses(
                'acbd18db4cc2f85cedef654fccc4a4d8+3',
                500, 200, 200) as req_mock:
            self.keep_client.put('foo', num_retries=1, copies=2)
        self.assertEqual(3, req_mock.call_count)

    def test_fail_after_final_error(self):
        # First retry loop gets a 200 (can't achieve replication by
        # storing again on that server) and a 400 (can't retry that
        # server at all), so we shouldn't try a third request.
        with tutil.mock_keep_responses(
                'acbd18db4cc2f85cedef654fccc4a4d8+3',
                200, 400, 200) as req_mock:
            with self.assertRaises(arvados.errors.KeepWriteError):
                self.keep_client.put('foo', num_retries=1, copies=2)
        self.assertEqual(2, req_mock.call_count)


@parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
class KeepClientAPIErrorTest(unittest.TestCase, DiskCacheBase):
    disk_cache = False

    def tearDown(self):
        DiskCacheBase.tearDown(self)

    def test_api_fail(self):
        class ApiMock(object):
            def __getattr__(self, r):
                if r == "api_token":
                    return "abc"
                elif r == "insecure":
                    return False
                elif r == "config":
                    return lambda: {}
                else:
                    raise arvados.errors.KeepReadError()
        keep_client = arvados.KeepClient(api_client=ApiMock(),
                                         proxy='', local_store='',
                                         block_cache=self.make_block_cache(self.disk_cache))

        # The bug this is testing for is that if an API (not
        # keepstore) exception is thrown as part of a get(), the next
        # attempt to get that same block will result in a deadlock.
        # This is why there are two get()s in a row.  Unfortunately,
        # the failure mode for this test is that the test suite
        # deadlocks, there isn't a good way to avoid that without
        # adding a special case that has no use except for this test.

        with self.assertRaises(arvados.errors.KeepReadError):
            keep_client.get("acbd18db4cc2f85cedef654fccc4a4d8+3")
        with self.assertRaises(arvados.errors.KeepReadError):
            keep_client.get("acbd18db4cc2f85cedef654fccc4a4d8+3")


class KeepDiskCacheTestCase(unittest.TestCase, tutil.ApiClientMock):
    def setUp(self):
        self.api_client = self.mock_keep_services(count=2)
        self.data = b'xyzzy'
        self.locator = '1271ed5ef305aadabc605b1609e24c52'
        self.disk_cache_dir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.disk_cache_dir)

    @mock.patch('arvados.util._BaseDirectories.storage_path')
    def test_default_disk_cache_dir(self, storage_path):
        expected = Path(self.disk_cache_dir)
        storage_path.return_value = expected
        cache = arvados.keep.KeepBlockCache(disk_cache=True)
        storage_path.assert_called_with('keep')
        self.assertEqual(cache._disk_cache_dir, str(expected))

    @mock.patch('arvados.KeepClient.KeepService.get')
    def test_disk_cache_read(self, get_mock):
        # confirm it finds an existing cache block when the cache is
        # initialized.

        os.makedirs(os.path.join(self.disk_cache_dir, self.locator[0:3]))
        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock"), "wb") as f:
            f.write(self.data)

        # block cache should have found the existing block
        block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                  disk_cache_dir=self.disk_cache_dir)
        keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=block_cache)

        self.assertTrue(tutil.binary_compare(keep_client.get(self.locator), self.data))

        get_mock.assert_not_called()

    @mock.patch('arvados.KeepClient.KeepService.get')
    def test_disk_cache_share(self, get_mock):
        # confirm it finds a cache block written after the disk cache
        # was initialized.

        block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                  disk_cache_dir=self.disk_cache_dir)
        keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=block_cache)

        os.makedirs(os.path.join(self.disk_cache_dir, self.locator[0:3]))
        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock"), "wb") as f:
            f.write(self.data)

        # when we try to get the block, it'll check the disk and find it.
        self.assertTrue(tutil.binary_compare(keep_client.get(self.locator), self.data))

        get_mock.assert_not_called()

    def test_disk_cache_write(self):
        # confirm the cache block was created

        block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                  disk_cache_dir=self.disk_cache_dir)
        keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=block_cache)

        with tutil.mock_keep_responses(self.data, 200) as mock:
            self.assertTrue(tutil.binary_compare(keep_client.get(self.locator), self.data))

        self.assertIsNotNone(keep_client.get_from_cache(self.locator))

        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock"), "rb") as f:
            self.assertTrue(tutil.binary_compare(f.read(), self.data))

    def test_disk_cache_clean(self):
        # confirm that a tmp file in the cache is cleaned up

        os.makedirs(os.path.join(self.disk_cache_dir, self.locator[0:3]))
        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC.keepcacheblock"), "wb") as f:
            f.write(b"abc1")

        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC"), "wb") as f:
            f.write(b"abc2")

        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], "XYZABC"), "wb") as f:
            f.write(b"abc3")

        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC.keepcacheblock")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], "XYZABC")))

        block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                  disk_cache_dir=self.disk_cache_dir)

        # The tmp still hasn't been deleted because it was created in the last 60 seconds
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC.keepcacheblock")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], "XYZABC")))

        # Set the mtime to 61s in the past
        os.utime(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC.keepcacheblock"), times=(time.time()-61, time.time()-61))
        os.utime(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC"), times=(time.time()-61, time.time()-61))
        os.utime(os.path.join(self.disk_cache_dir, self.locator[0:3], "XYZABC"), times=(time.time()-61, time.time()-61))

        block_cache2 = arvados.keep.KeepBlockCache(disk_cache=True,
                                                   disk_cache_dir=self.disk_cache_dir)

        # Tmp should be gone but the other ones are safe.
        self.assertFalse(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC.keepcacheblock")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], "tmpXYZABC")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], "XYZABC")))

    @mock.patch('arvados.KeepClient.KeepService.get')
    def test_disk_cache_cap(self, get_mock):
        # confirm that the cache is kept to the desired limit

        os.makedirs(os.path.join(self.disk_cache_dir, self.locator[0:3]))
        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock"), "wb") as f:
            f.write(self.data)

        os.makedirs(os.path.join(self.disk_cache_dir, "acb"))
        with open(os.path.join(self.disk_cache_dir, "acb", "acbd18db4cc2f85cedef654fccc4a4d8.keepcacheblock"), "wb") as f:
            f.write(b"foo")

        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, "acb", "acbd18db4cc2f85cedef654fccc4a4d8.keepcacheblock")))

        block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                   disk_cache_dir=self.disk_cache_dir,
                                                   max_slots=1)

        self.assertFalse(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, "acb", "acbd18db4cc2f85cedef654fccc4a4d8.keepcacheblock")))

    @mock.patch('arvados.KeepClient.KeepService.get')
    def test_disk_cache_share(self, get_mock):
        # confirm that a second cache doesn't delete files that belong to the first cache.

        os.makedirs(os.path.join(self.disk_cache_dir, self.locator[0:3]))
        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock"), "wb") as f:
            f.write(self.data)

        os.makedirs(os.path.join(self.disk_cache_dir, "acb"))
        with open(os.path.join(self.disk_cache_dir, "acb", "acbd18db4cc2f85cedef654fccc4a4d8.keepcacheblock"), "wb") as f:
            f.write(b"foo")

        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, "acb", "acbd18db4cc2f85cedef654fccc4a4d8.keepcacheblock")))

        block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                   disk_cache_dir=self.disk_cache_dir,
                                                   max_slots=2)

        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, "acb", "acbd18db4cc2f85cedef654fccc4a4d8.keepcacheblock")))

        block_cache2 = arvados.keep.KeepBlockCache(disk_cache=True,
                                                   disk_cache_dir=self.disk_cache_dir,
                                                   max_slots=1)

        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock")))
        self.assertTrue(os.path.exists(os.path.join(self.disk_cache_dir, "acb", "acbd18db4cc2f85cedef654fccc4a4d8.keepcacheblock")))

    def test_disk_cache_error(self):
        os.chmod(self.disk_cache_dir, stat.S_IRUSR)

        # Fail during cache initialization.
        with self.assertRaises(OSError):
            block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                      disk_cache_dir=self.disk_cache_dir)

    def test_disk_cache_write_error(self):
        block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                  disk_cache_dir=self.disk_cache_dir)

        keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=block_cache)

        # Make the cache dir read-only
        os.makedirs(os.path.join(self.disk_cache_dir, self.locator[0:3]))
        os.chmod(os.path.join(self.disk_cache_dir, self.locator[0:3]), stat.S_IRUSR)

        # Cache fails
        with self.assertRaises(arvados.errors.KeepCacheError):
            with tutil.mock_keep_responses(self.data, 200) as mock:
                keep_client.get(self.locator)

    def test_disk_cache_retry_write_error(self):
        block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                  disk_cache_dir=self.disk_cache_dir)

        keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=block_cache)

        called = False
        realmmap = mmap.mmap
        def sideeffect_mmap(*args, **kwargs):
            nonlocal called
            if not called:
                called = True
                raise OSError(errno.ENOSPC, "no space")
            else:
                return realmmap(*args, **kwargs)

        with patch('mmap.mmap') as mockmmap:
            mockmmap.side_effect = sideeffect_mmap

            cache_max_before = block_cache.cache_max

            with tutil.mock_keep_responses(self.data, 200) as mock:
                self.assertTrue(tutil.binary_compare(keep_client.get(self.locator), self.data))

            self.assertIsNotNone(keep_client.get_from_cache(self.locator))

        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock"), "rb") as f:
            self.assertTrue(tutil.binary_compare(f.read(), self.data))

        # shrank the cache in response to ENOSPC
        self.assertTrue(cache_max_before > block_cache.cache_max)

    def test_disk_cache_retry_write_error2(self):
        block_cache = arvados.keep.KeepBlockCache(disk_cache=True,
                                                  disk_cache_dir=self.disk_cache_dir)

        keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=block_cache)

        called = False
        realmmap = mmap.mmap
        def sideeffect_mmap(*args, **kwargs):
            nonlocal called
            if not called:
                called = True
                raise OSError(errno.ENOMEM, "no memory")
            else:
                return realmmap(*args, **kwargs)

        with patch('mmap.mmap') as mockmmap:
            mockmmap.side_effect = sideeffect_mmap

            slots_before = block_cache._max_slots

            with tutil.mock_keep_responses(self.data, 200) as mock:
                self.assertTrue(tutil.binary_compare(keep_client.get(self.locator), self.data))

            self.assertIsNotNone(keep_client.get_from_cache(self.locator))

        with open(os.path.join(self.disk_cache_dir, self.locator[0:3], self.locator+".keepcacheblock"), "rb") as f:
            self.assertTrue(tutil.binary_compare(f.read(), self.data))

        # shrank the cache in response to ENOMEM
        self.assertTrue(slots_before > block_cache._max_slots)