#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import inspect
import time
import unittest

import mock
import pytest
from mox import IgnoreArg, IsA, Mox
from requests.auth import AuthBase
from thrift.transport import TTransport
from twitter.common.quantity import Amount, Time
from twitter.common.zookeeper.kazoo_client import TwitterKazooClient
from twitter.common.zookeeper.serverset.endpoint import ServiceInstance

import apache.aurora.client.api.scheduler_client as scheduler_client
from apache.aurora.common.auth.auth_module import AuthModule
from apache.aurora.common.cluster import Cluster
from apache.aurora.common.transport import TRequestsTransport

import gen.apache.aurora.api.AuroraAdmin as AuroraAdmin
import gen.apache.aurora.api.AuroraSchedulerManager as AuroraSchedulerManager
from gen.apache.aurora.api.constants import BYPASS_LEADER_REDIRECT_HEADER_NAME
from gen.apache.aurora.api.ttypes import (
    ExplicitReconciliationSettings,
    Hosts,
    JobConfiguration,
    JobKey,
    JobUpdateQuery,
    JobUpdateRequest,
    ResourceAggregate,
    Response,
    ResponseCode,
    ResponseDetail,
    ScheduleStatus,
    SlaPolicy,
    TaskQuery
)

ROLE = 'foorole'
JOB_NAME = 'barjobname'
JOB_ENV = 'devel'
JOB_KEY = JobKey(role=ROLE, environment=JOB_ENV, name=JOB_NAME)
DEFAULT_RESPONSE = Response()


def test_coverage():
  """Make sure a new thrift RPC doesn't get added without minimal test coverage."""
  for name, klass in inspect.getmembers(AuroraAdmin) + inspect.getmembers(AuroraSchedulerManager):
    if name.endswith('_args'):
      rpc_name = name[:-len('_args')]
      assert hasattr(TestSchedulerProxyAdminInjection, 'test_%s' % rpc_name), (
          'No test defined for RPC %s' % rpc_name)


class TestSchedulerProxyInjection(unittest.TestCase):
  def setUp(self):
    self.mox = Mox()

    self.mox.StubOutClassWithMocks(AuroraAdmin, 'Client')
    self.mox.StubOutClassWithMocks(scheduler_client, 'SchedulerClient')

    self.mock_scheduler_client = self.mox.CreateMock(scheduler_client.SchedulerClient)
    self.mock_thrift_client = self.mox.CreateMock(AuroraAdmin.Client)

    scheduler_client.SchedulerClient.get(IgnoreArg(), verbose=IgnoreArg()).AndReturn(
        self.mock_scheduler_client)
    self.mock_scheduler_client.get_thrift_client().AndReturn(self.mock_thrift_client)

  def tearDown(self):
    self.mox.UnsetStubs()
    self.mox.VerifyAll()

  def make_scheduler_proxy(self):
    return scheduler_client.SchedulerProxy(Cluster(name='local'))

  def test_startCronJob(self):
    self.mock_thrift_client.startCronJob(IsA(JobKey)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().startCronJob(JOB_KEY)

  def test_createJob(self):
    self.mock_thrift_client.createJob(
        IsA(JobConfiguration)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().createJob(JobConfiguration())

  def test_replaceCronTemplate(self):
    self.mock_thrift_client.replaceCronTemplate(
        IsA(JobConfiguration)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().replaceCronTemplate(JobConfiguration())

  def test_scheduleCronJob(self):
    self.mock_thrift_client.scheduleCronJob(
        IsA(JobConfiguration)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().scheduleCronJob(JobConfiguration())

  def test_descheduleCronJob(self):
    self.mock_thrift_client.descheduleCronJob(
        IsA(JobKey)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().descheduleCronJob(JOB_KEY)

  def test_populateJobConfig(self):
    self.mock_thrift_client.populateJobConfig(IsA(JobConfiguration)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().populateJobConfig(JobConfiguration())

  def test_restartShards(self):
    self.mock_thrift_client.restartShards(IsA(JobKey), IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().restartShards(JOB_KEY, {0})

  def test_getTasksStatus(self):
    self.mock_thrift_client.getTasksStatus(IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().getTasksStatus(TaskQuery())

  def test_getJobs(self):
    self.mock_thrift_client.getJobs(IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().getJobs(ROLE)

  def test_killTasks(self):
    self.mock_thrift_client.killTasks(IsA(JobKey), IgnoreArg(), None).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().killTasks(JobKey(), {0}, None)

  def test_getQuota(self):
    self.mock_thrift_client.getQuota(IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().getQuota(ROLE)

  def test_addInstances(self):
    self.mock_thrift_client.addInstances(IsA(JobKey), 1).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().addInstances(JobKey(), 1)

  def test_getJobUpdateSummaries(self):
    self.mock_thrift_client.getJobUpdateSummaries(IsA(JobUpdateQuery)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().getJobUpdateSummaries(JobUpdateQuery())

  def test_getJobUpdateDetails(self):
    self.mock_thrift_client.getJobUpdateDetails('update_id').AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().getJobUpdateDetails('update_id')

  def test_startJobUpdate(self):
    self.mock_thrift_client.startJobUpdate(
        IsA(JobUpdateRequest)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().startJobUpdate(JobUpdateRequest())

  def test_pauseJobUpdate(self):
    self.mock_thrift_client.pauseJobUpdate('update_id').AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().pauseJobUpdate('update_id')

  def test_resumeJobUpdate(self):
    self.mock_thrift_client.resumeJobUpdate(
        'update_id').AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().resumeJobUpdate('update_id')

  def test_abortJobUpdate(self):
    self.mock_thrift_client.abortJobUpdate('update_id').AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().abortJobUpdate('update_id')

  def test_rollbackJobUpdate(self):
    self.mock_thrift_client.rollbackJobUpdate('update_id').AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().rollbackJobUpdate('update_id')

  def test_pulseJobUpdate(self):
    self.mock_thrift_client.pulseJobUpdate('update_id').AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().pulseJobUpdate('update_id')

  def test_raise_auth_error(self):
    self.mock_thrift_client.killTasks(None, None, None).AndRaise(TRequestsTransport.AuthError())
    self.mock_scheduler_client.get_failed_auth_message().AndReturn('failed auth')
    self.mox.ReplayAll()
    with pytest.raises(scheduler_client.SchedulerProxy.AuthError):
      self.make_scheduler_proxy().killTasks(None, None, None)


class TestSchedulerProxyAdminInjection(TestSchedulerProxyInjection):
  def test_startMaintenance(self):
    self.mock_thrift_client.startMaintenance(
      IsA(Hosts)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().startMaintenance(Hosts())

  def test_drainHosts(self):
    self.mock_thrift_client.drainHosts(IsA(Hosts)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().drainHosts(Hosts())

  def test_slaDrainHosts(self):
    self.mock_thrift_client.slaDrainHosts(
      IsA(Hosts),
      IsA(SlaPolicy),
      10).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().slaDrainHosts(Hosts(), SlaPolicy(), 10)

  def test_maintenanceStatus(self):
    self.mock_thrift_client.maintenanceStatus(
      IsA(Hosts)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().maintenanceStatus(Hosts())

  def test_endMaintenance(self):
    self.mock_thrift_client.endMaintenance(IsA(Hosts)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().endMaintenance(Hosts())

  def test_setQuota(self):
    self.mock_thrift_client.setQuota(
        IgnoreArg(),
        IsA(ResourceAggregate)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().setQuota(ROLE, ResourceAggregate())

  def test_forceTaskState(self):
    self.mock_thrift_client.forceTaskState(
        'taskid',
        IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().forceTaskState('taskid', ScheduleStatus.LOST)

  def test_performBackup(self):
    self.mock_thrift_client.performBackup().AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().performBackup()

  def test_listBackups(self):
    self.mock_thrift_client.listBackups().AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().listBackups()

  def test_stageRecovery(self):
    self.mock_thrift_client.stageRecovery(
        IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().stageRecovery(TaskQuery())

  def test_queryRecovery(self):
    self.mock_thrift_client.queryRecovery(
      IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().queryRecovery(TaskQuery())

  def test_deleteRecoveryTasks(self):
    self.mock_thrift_client.deleteRecoveryTasks(
        IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().deleteRecoveryTasks(TaskQuery())

  def test_commitRecovery(self):
    self.mock_thrift_client.commitRecovery().AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().commitRecovery()

  def test_unloadRecovery(self):
    self.mock_thrift_client.unloadRecovery().AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().unloadRecovery()

  def test_snapshot(self):
    self.mock_thrift_client.snapshot().AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().snapshot()

  def test_pruneTasks(self):
    t = TaskQuery()
    self.mock_thrift_client.pruneTasks(IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().pruneTasks(t)

  def test_triggerExplicitTaskReconciliation(self):
    self.mock_thrift_client.triggerExplicitTaskReconciliation(
      IsA(ExplicitReconciliationSettings)).AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().triggerExplicitTaskReconciliation(
      ExplicitReconciliationSettings(batchSize=None))

  def test_triggerImplicitTaskReconciliation(self):
    self.mock_thrift_client.triggerImplicitTaskReconciliation().AndReturn(DEFAULT_RESPONSE)
    self.mox.ReplayAll()
    self.make_scheduler_proxy().triggerImplicitTaskReconciliation()


def mock_auth():
  auth_mock = mock.create_autospec(spec=AuthModule, instance=True)
  auth_mock.auth.return_value = mock.create_autospec(AuthBase)
  return auth_mock


@pytest.mark.parametrize('scheme', ('http', 'https'))
def test_url_when_not_connected_and_cluster_has_no_proxy_url(scheme):
  host = 'some-host.example.com'
  port = 31181

  mock_zk = mock.create_autospec(spec=TwitterKazooClient, instance=True)

  service_json = '''{
    "additionalEndpoints": {
        "%(scheme)s": {
            "host": "%(host)s",
            "port": %(port)d
        }
    },
    "serviceEndpoint": {
        "host": "%(host)s",
        "port": %(port)d
    },
    "shard": 0,
    "status": "ALIVE"
  }''' % dict(host=host, port=port, scheme=scheme)

  service_endpoints = [ServiceInstance.unpack(service_json)]

  def make_mock_client(proxy_url):
    client = scheduler_client.ZookeeperSchedulerClient(
        Cluster(proxy_url=proxy_url),
        auth=None,
        user_agent='Some-User-Agent',
        _deadline=lambda x, **kws: x())
    client.get_scheduler_serverset = mock.MagicMock(return_value=(mock_zk, service_endpoints))
    client.SERVERSET_TIMEOUT = Amount(0, Time.SECONDS)
    client._connect_scheduler = mock.MagicMock()
    return client

  client = make_mock_client(proxy_url=None)
  assert client.url == '%s://%s:%d' % (scheme, host, port)
  assert client.url == client.raw_url
  client._connect_scheduler.assert_has_calls([])

  client = make_mock_client(proxy_url='https://scheduler.proxy')
  assert client.url == 'https://scheduler.proxy'
  assert client.raw_url == '%s://%s:%d' % (scheme, host, port)
  client._connect_scheduler.assert_has_calls([])

  client = make_mock_client(proxy_url=None)
  client.get_thrift_client()
  assert client.url == '%s://%s:%d' % (scheme, host, port)
  client._connect_scheduler.assert_has_calls([mock.call('%s://%s:%d/api' % (scheme, host, port))])
  client._connect_scheduler.reset_mock()
  client.get_thrift_client()
  client._connect_scheduler.assert_has_calls([])


class TestSchedulerClient(unittest.TestCase):

  @mock.patch('apache.aurora.client.api.scheduler_client.TRequestsTransport',
              spec=TRequestsTransport)
  def test_connect_scheduler(self, mock_client):
    mock_client.return_value.open.side_effect = [TTransport.TTransportException, True]
    mock_time = mock.create_autospec(spec=time, instance=True)

    client = scheduler_client.SchedulerClient(mock_auth(), 'Some-User-Agent', verbose=True)
    client._connect_scheduler('https://scheduler.example.com:1337', mock_time)

    assert mock_client.return_value.open.has_calls(mock.call(), mock.call())
    mock_time.sleep.assert_called_once_with(
        scheduler_client.SchedulerClient.RETRY_TIMEOUT.as_(Time.SECONDS))

  @mock.patch('apache.aurora.client.api.scheduler_client.TRequestsTransport',
              spec=TRequestsTransport)
  def test_connect_scheduler_with_user_agent(self, mock_transport):
    mock_transport.return_value.open.side_effect = [TTransport.TTransportException, True]
    mock_time = mock.create_autospec(spec=time, instance=True)

    auth = mock_auth()
    user_agent = 'Some-User-Agent'

    client = scheduler_client.SchedulerClient(auth, user_agent, verbose=True)

    uri = 'https://scheduler.example.com:1337'
    client._connect_scheduler(uri, mock_time)

    mock_transport.assert_called_once_with(
        uri,
        auth=auth.auth(),
        user_agent=user_agent,
        session_factory=mock.ANY)

  @mock.patch('apache.aurora.client.api.scheduler_client.TRequestsTransport',
              spec=TRequestsTransport)
  def test_connect_scheduler_without_bypass_leader_redirect(self, mock_transport):
    mock_transport.return_value.open.side_effect = [TTransport.TTransportException, True]
    mock_time = mock.create_autospec(spec=time, instance=True)

    auth = mock_auth()
    user_agent = 'Some-User-Agent'

    client = scheduler_client.SchedulerClient(
        auth,
        user_agent,
        verbose=True,
        bypass_leader_redirect=False)

    uri = 'https://scheduler.example.com:1337'
    client._connect_scheduler(uri, mock_time)

    mock_transport.assert_called_once_with(
        uri,
        auth=auth.auth(),
        user_agent=user_agent,
        session_factory=mock.ANY)

    _, _, kwargs = mock_transport.mock_calls[0]
    session = kwargs['session_factory']()
    assert session.headers.get(BYPASS_LEADER_REDIRECT_HEADER_NAME) is None

  @mock.patch('apache.aurora.client.api.scheduler_client.TRequestsTransport',
              spec=TRequestsTransport)
  def test_connect_scheduler_with_bypass_leader_redirect(self, mock_transport):
    mock_transport.return_value.open.side_effect = [TTransport.TTransportException, True]
    mock_time = mock.create_autospec(spec=time, instance=True)

    auth = mock_auth()
    user_agent = 'Some-User-Agent'

    client = scheduler_client.SchedulerClient(
        auth,
        user_agent,
        verbose=True,
        bypass_leader_redirect=True)

    uri = 'https://scheduler.example.com:1337'
    client._connect_scheduler(uri, mock_time)

    mock_transport.assert_called_once_with(
        uri,
        auth=auth.auth(),
        user_agent=user_agent,
        session_factory=mock.ANY)

    _, _, kwargs = mock_transport.mock_calls[0]
    session = kwargs['session_factory']()
    assert session.headers[BYPASS_LEADER_REDIRECT_HEADER_NAME] == 'true'

  @mock.patch('apache.aurora.client.api.scheduler_client.SchedulerClient',
              spec=scheduler_client.SchedulerClient)
  @mock.patch('threading._Event.wait')
  def test_transient_error(self, _, client):
    mock_scheduler_client = mock.create_autospec(
        spec=scheduler_client.SchedulerClient,
        spec_set=False,
        instance=True)
    mock_thrift_client = mock.create_autospec(spec=AuroraAdmin.Client, instance=True)
    mock_thrift_client.killTasks.side_effect = [
        Response(responseCode=ResponseCode.ERROR_TRANSIENT,
                 details=[ResponseDetail(message="message1"), ResponseDetail(message="message2")]),
        Response(responseCode=ResponseCode.ERROR_TRANSIENT),
        Response(responseCode=ResponseCode.OK)]

    mock_scheduler_client.get_thrift_client.return_value = mock_thrift_client
    client.get.return_value = mock_scheduler_client

    proxy = scheduler_client.SchedulerProxy(Cluster(name='local'))
    proxy.killTasks(JobKey(), None, None)

    assert mock_thrift_client.killTasks.call_count == 3

  @mock.patch('apache.aurora.client.api.scheduler_client.SchedulerClient',
              spec=scheduler_client.SchedulerClient)
  @mock.patch('threading._Event.wait')
  def test_performBackup_retriable_errors(self, mock_wait, mock_client):
    mock_scheduler_client = mock.create_autospec(
        spec=scheduler_client.SchedulerClient,
        spec_set=False,
        instance=True)
    mock_thrift_client = mock.create_autospec(spec=AuroraAdmin.Client, instance=True)
    mock_thrift_client.performBackup.side_effect = [
      Response(responseCode=ResponseCode.ERROR_TRANSIENT),
      scheduler_client.SchedulerProxy.TimeoutError,
      Response(responseCode=ResponseCode.OK)]

    mock_scheduler_client.get_thrift_client.return_value = mock_thrift_client
    mock_client.get.return_value = mock_scheduler_client

    proxy = scheduler_client.SchedulerProxy(Cluster(name='local'))
    proxy.performBackup()

    assert mock_thrift_client.performBackup.call_count == 3
    assert mock_wait.call_count == 2

  @mock.patch('apache.aurora.client.api.scheduler_client.SchedulerClient',
              spec=scheduler_client.SchedulerClient)
  @mock.patch('threading._Event.wait')
  def test_performBackup_transport_exception(self, mock_wait, mock_client):
    mock_scheduler_client = mock.create_autospec(
      spec=scheduler_client.SchedulerClient,
      spec_set=False,
      instance=True)

    mock_thrift_client = mock.create_autospec(spec=AuroraAdmin.Client, instance=True)
    mock_thrift_client.performBackup.side_effect = TTransport.TTransportException('error')
    mock_scheduler_client.get_thrift_client.return_value = mock_thrift_client
    mock_client.get.return_value = mock_scheduler_client

    proxy = scheduler_client.SchedulerProxy(Cluster(name='local'))
    with pytest.raises(scheduler_client.SchedulerProxy.NotRetriableError):
      proxy.performBackup()

    assert mock_thrift_client.performBackup.call_count == 1
    assert not mock_wait.called

  @mock.patch('apache.aurora.client.api.scheduler_client.SchedulerClient',
              spec=scheduler_client.SchedulerClient)
  @mock.patch('threading._Event.wait')
  def test_getTierConfigs_transport_exception(self, mock_wait, mock_client):
    mock_scheduler_client = mock.create_autospec(
      spec=scheduler_client.SchedulerClient,
      spec_set=False,
      instance=True)

    mock_thrift_client = mock.create_autospec(spec=AuroraAdmin.Client, instance=True)
    mock_thrift_client.getTierConfigs.side_effect = [
      TTransport.TTransportException('error'),
      Response(responseCode=ResponseCode.OK)
    ]
    mock_scheduler_client.get_thrift_client.return_value = mock_thrift_client
    mock_client.get.return_value = mock_scheduler_client

    proxy = scheduler_client.SchedulerProxy(Cluster(name='local'))
    proxy.getTierConfigs(retry=True)

    assert mock_thrift_client.getTierConfigs.call_count == 2
    assert mock_wait.call_count == 1

  @mock.patch('apache.aurora.client.api.scheduler_client.SchedulerClient',
              spec=scheduler_client.SchedulerClient)
  def test_unknown_connection_error(self, client):
    mock_scheduler_client = mock.create_autospec(spec=scheduler_client.SchedulerClient,
                                                 instance=True)
    client.get.return_value = mock_scheduler_client
    proxy = scheduler_client.SchedulerProxy(Cluster(name='local'))

    # unknown, transient connection error
    mock_scheduler_client.get_thrift_client.side_effect = RuntimeError
    with pytest.raises(Exception):
      proxy.client()

    # successful connection on re-attempt
    mock_scheduler_client.get_thrift_client.side_effect = None
    assert proxy.client() is not None

  @mock.patch('apache.aurora.client.api.scheduler_client.TRequestsTransport',
              spec=TRequestsTransport)
  def test_connect_direct_scheduler_with_user_agent(self, mock_transport):
    mock_transport.return_value.open.side_effect = [TTransport.TTransportException, True]
    mock_time = mock.create_autospec(spec=time, instance=True)

    auth = mock_auth()
    user_agent = 'Some-User-Agent'
    uri = 'https://scheduler.example.com:1337'

    client = scheduler_client.DirectSchedulerClient(
        uri,
        auth=auth,
        verbose=True,
        user_agent=user_agent)

    client._connect_scheduler(uri, mock_time)

    mock_transport.assert_called_once_with(
        uri,
        auth=auth.auth(),
        user_agent=user_agent,
        session_factory=mock.ANY)

  @mock.patch('apache.aurora.client.api.scheduler_client.TRequestsTransport',
              spec=TRequestsTransport)
  def test_connect_zookeeper_client_with_auth(self, mock_transport):
    mock_transport.return_value.open.side_effect = [TTransport.TTransportException, True]
    mock_time = mock.create_autospec(spec=time, instance=True)

    user_agent = 'Some-User-Agent'
    uri = 'https://scheduler.example.com:1337'
    auth = mock_auth()
    cluster = Cluster(zk='zk', zk_port='2181')

    def auth_factory(_):
      return auth

    client = scheduler_client.SchedulerClient.get(
        cluster,
        auth_factory=auth_factory,
        user_agent=user_agent)

    client._connect_scheduler(uri, mock_time)

    mock_transport.assert_called_once_with(
        uri,
        auth=auth.auth(),
        user_agent=user_agent,
        session_factory=mock.ANY)

  @mock.patch('apache.aurora.client.api.scheduler_client.TRequestsTransport',
              spec=TRequestsTransport)
  def test_connect_direct_client_with_auth(self, mock_transport):
    mock_transport.return_value.open.side_effect = [TTransport.TTransportException, True]
    mock_time = mock.create_autospec(spec=time, instance=True)

    user_agent = 'Some-User-Agent'
    uri = 'https://scheduler.example.com:1337'
    auth = mock_auth()
    cluster = Cluster(scheduler_uri='uri')

    def auth_factory(_):
      return auth

    client = scheduler_client.SchedulerClient.get(
        cluster,
        auth_factory=auth_factory,
        user_agent=user_agent)

    client._connect_scheduler(uri, mock_time)

    mock_transport.assert_called_once_with(
        uri,
        auth=auth.auth(),
        user_agent=user_agent,
        session_factory=mock.ANY)

  def test_no_zk_or_scheduler_uri(self):
    cluster = None
    with self.assertRaises(TypeError):
      scheduler_client.SchedulerClient.get(cluster)
    cluster = Cluster()
    with self.assertRaises(ValueError):
      scheduler_client.SchedulerClient.get(cluster)

  def test__internal_connect(self):
    client = scheduler_client.SchedulerClient(mock_auth(), 'Some-User-Agent', verbose=True)
    self.assertIsNone(client._connect())