path: root/test/test_coordinator.py

# pylint: skip-file
from __future__ import absolute_import
import time

import pytest

from kafka.client_async import KafkaClient
from kafka.consumer.subscription_state import (
    SubscriptionState, ConsumerRebalanceListener)
from kafka.coordinator.assignors.range import RangePartitionAssignor
from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
from kafka.coordinator.assignors.sticky.sticky_assignor import StickyPartitionAssignor
from kafka.coordinator.base import Generation, MemberState, HeartbeatThread
from kafka.coordinator.consumer import ConsumerCoordinator
from kafka.coordinator.protocol import (
    ConsumerProtocolMemberMetadata, ConsumerProtocolMemberAssignment)
import kafka.errors as Errors
from kafka.future import Future
from kafka.metrics import Metrics
from kafka.protocol.commit import (
    OffsetCommitRequest, OffsetCommitResponse,
    OffsetFetchRequest, OffsetFetchResponse)
from kafka.protocol.metadata import MetadataResponse
from kafka.structs import OffsetAndMetadata, TopicPartition
from kafka.util import WeakMethod


@pytest.fixture
def client(conn):
    return KafkaClient(api_version=(0, 9))

@pytest.fixture
def coordinator(client):
    return ConsumerCoordinator(client, SubscriptionState(), Metrics())


def test_init(client, coordinator):
    # metadata update on init
    assert client.cluster._need_update is True
    assert WeakMethod(coordinator._handle_metadata_update) in client.cluster._listeners


@pytest.mark.parametrize("api_version", [(0, 8, 0), (0, 8, 1), (0, 8, 2), (0, 9)])
def test_autocommit_enable_api_version(client, api_version):
    coordinator = ConsumerCoordinator(client, SubscriptionState(),
                                      Metrics(),
                                      enable_auto_commit=True,
                                      session_timeout_ms=30000,   # session_timeout_ms and max_poll_interval_ms
                                      max_poll_interval_ms=30000, # should be the same to avoid KafkaConfigurationError
                                      group_id='foobar',
                                      api_version=api_version)
    if api_version < (0, 8, 1):
        assert coordinator.config['enable_auto_commit'] is False
    else:
        assert coordinator.config['enable_auto_commit'] is True


def test_protocol_type(coordinator):
    assert coordinator.protocol_type() == 'consumer'


def test_group_protocols(coordinator):
    # Requires a subscription
    try:
        coordinator.group_protocols()
    except Errors.IllegalStateError:
        pass
    else:
        assert False, 'Exception not raised when expected'

    coordinator._subscription.subscribe(topics=['foobar'])
    assert coordinator.group_protocols() == [
        ('range', ConsumerProtocolMemberMetadata(
            RangePartitionAssignor.version,
            ['foobar'],
            b'')),
        ('roundrobin', ConsumerProtocolMemberMetadata(
            RoundRobinPartitionAssignor.version,
            ['foobar'],
            b'')),
        ('sticky', ConsumerProtocolMemberMetadata(
            StickyPartitionAssignor.version,
            ['foobar'],
            b'')),
    ]


@pytest.mark.parametrize('api_version', [(0, 8, 0), (0, 8, 1), (0, 8, 2), (0, 9)])
def test_pattern_subscription(coordinator, api_version):
    coordinator.config['api_version'] = api_version
    coordinator._subscription.subscribe(pattern='foo')
    assert coordinator._subscription.subscription == set([])
    assert coordinator._metadata_snapshot == coordinator._build_metadata_snapshot(coordinator._subscription, {})

    cluster = coordinator._client.cluster
    cluster.update_metadata(MetadataResponse[0](
        # brokers
        [(0, 'foo', 12), (1, 'bar', 34)],
        # topics
        [(0, 'fizz', []),
         (0, 'foo1', [(0, 0, 0, [], [])]),
         (0, 'foo2', [(0, 0, 1, [], [])])]))
    assert coordinator._subscription.subscription == {'foo1', 'foo2'}

    # 0.9 consumers should trigger dynamic partition assignment
    if api_version >= (0, 9):
        assert coordinator._subscription.assignment == {}

    # earlier consumers get all partitions assigned locally
    else:
        assert set(coordinator._subscription.assignment.keys()) == {TopicPartition('foo1', 0),
                                                                    TopicPartition('foo2', 0)}


def test_lookup_assignor(coordinator):
    assert coordinator._lookup_assignor('roundrobin') is RoundRobinPartitionAssignor
    assert coordinator._lookup_assignor('range') is RangePartitionAssignor
    assert coordinator._lookup_assignor('sticky') is StickyPartitionAssignor
    assert coordinator._lookup_assignor('foobar') is None


def test_join_complete(mocker, coordinator):
    coordinator._subscription.subscribe(topics=['foobar'])
    assignor = RoundRobinPartitionAssignor()
    coordinator.config['assignors'] = (assignor,)
    mocker.spy(assignor, 'on_assignment')
    assert assignor.on_assignment.call_count == 0
    assignment = ConsumerProtocolMemberAssignment(0, [('foobar', [0, 1])], b'')
    coordinator._on_join_complete(0, 'member-foo', 'roundrobin', assignment.encode())
    assert assignor.on_assignment.call_count == 1
    assignor.on_assignment.assert_called_with(assignment)


def test_join_complete_with_sticky_assignor(mocker, coordinator):
    coordinator._subscription.subscribe(topics=['foobar'])
    assignor = StickyPartitionAssignor()
    coordinator.config['assignors'] = (assignor,)
    mocker.spy(assignor, 'on_assignment')
    mocker.spy(assignor, 'on_generation_assignment')
    assert assignor.on_assignment.call_count == 0
    assert assignor.on_generation_assignment.call_count == 0
    assignment = ConsumerProtocolMemberAssignment(0, [('foobar', [0, 1])], b'')
    coordinator._on_join_complete(0, 'member-foo', 'sticky', assignment.encode())
    assert assignor.on_assignment.call_count == 1
    assert assignor.on_generation_assignment.call_count == 1
    assignor.on_assignment.assert_called_with(assignment)
    assignor.on_generation_assignment.assert_called_with(0)


def test_subscription_listener(mocker, coordinator):
    listener = mocker.MagicMock(spec=ConsumerRebalanceListener)
    coordinator._subscription.subscribe(
        topics=['foobar'],
        listener=listener)

    coordinator._on_join_prepare(0, 'member-foo')
    assert listener.on_partitions_revoked.call_count == 1
    listener.on_partitions_revoked.assert_called_with(set([]))

    assignment = ConsumerProtocolMemberAssignment(0, [('foobar', [0, 1])], b'')
    coordinator._on_join_complete(
        0, 'member-foo', 'roundrobin', assignment.encode())
    assert listener.on_partitions_assigned.call_count == 1
    listener.on_partitions_assigned.assert_called_with({TopicPartition('foobar', 0), TopicPartition('foobar', 1)})


def test_subscription_listener_failure(mocker, coordinator):
    listener = mocker.MagicMock(spec=ConsumerRebalanceListener)
    coordinator._subscription.subscribe(
        topics=['foobar'],
        listener=listener)

    # exception raised in listener should not be re-raised by coordinator
    listener.on_partitions_revoked.side_effect = Exception('crash')
    coordinator._on_join_prepare(0, 'member-foo')
    assert listener.on_partitions_revoked.call_count == 1

    assignment = ConsumerProtocolMemberAssignment(0, [('foobar', [0, 1])], b'')
    coordinator._on_join_complete(
        0, 'member-foo', 'roundrobin', assignment.encode())
    assert listener.on_partitions_assigned.call_count == 1


def test_perform_assignment(mocker, coordinator):
    member_metadata = {
        'member-foo': ConsumerProtocolMemberMetadata(0, ['foo1'], b''),
        'member-bar': ConsumerProtocolMemberMetadata(0, ['foo1'], b'')
    }
    assignments = {
        'member-foo': ConsumerProtocolMemberAssignment(
            0, [('foo1', [0])], b''),
        'member-bar': ConsumerProtocolMemberAssignment(
            0, [('foo1', [1])], b'')
    }

    mocker.patch.object(RoundRobinPartitionAssignor, 'assign')
    RoundRobinPartitionAssignor.assign.return_value = assignments

    ret = coordinator._perform_assignment(
        'member-foo', 'roundrobin',
        [(member, metadata.encode())
         for member, metadata in member_metadata.items()])

    assert RoundRobinPartitionAssignor.assign.call_count == 1
    RoundRobinPartitionAssignor.assign.assert_called_with(
        coordinator._client.cluster, member_metadata)
    assert ret == assignments


def test_on_join_prepare(coordinator):
    coordinator._subscription.subscribe(topics=['foobar'])
    coordinator._on_join_prepare(0, 'member-foo')


def test_need_rejoin(coordinator):
    # No subscription - no rejoin
    assert coordinator.need_rejoin() is False

    coordinator._subscription.subscribe(topics=['foobar'])
    assert coordinator.need_rejoin() is True


def test_refresh_committed_offsets_if_needed(mocker, coordinator):
    mocker.patch.object(ConsumerCoordinator, 'fetch_committed_offsets',
                        return_value = {
                            TopicPartition('foobar', 0): OffsetAndMetadata(123, b''),
                            TopicPartition('foobar', 1): OffsetAndMetadata(234, b'')})
    coordinator._subscription.assign_from_user([TopicPartition('foobar', 0)])
    assert coordinator._subscription.needs_fetch_committed_offsets is True
    coordinator.refresh_committed_offsets_if_needed()
    assignment = coordinator._subscription.assignment
    assert assignment[TopicPartition('foobar', 0)].committed == OffsetAndMetadata(123, b'')
    assert TopicPartition('foobar', 1) not in assignment
    assert coordinator._subscription.needs_fetch_committed_offsets is False


def test_fetch_committed_offsets(mocker, coordinator):

    # No partitions, no IO polling
    mocker.patch.object(coordinator._client, 'poll')
    assert coordinator.fetch_committed_offsets([]) == {}
    assert coordinator._client.poll.call_count == 0

    # general case -- send offset fetch request, get successful future
    mocker.patch.object(coordinator, 'ensure_coordinator_ready')
    mocker.patch.object(coordinator, '_send_offset_fetch_request',
                        return_value=Future().success('foobar'))
    partitions = [TopicPartition('foobar', 0)]
    ret = coordinator.fetch_committed_offsets(partitions)
    assert ret == 'foobar'
    coordinator._send_offset_fetch_request.assert_called_with(partitions)
    assert coordinator._client.poll.call_count == 1

    # Failed future is raised if not retriable
    coordinator._send_offset_fetch_request.return_value = Future().failure(AssertionError)
    coordinator._client.poll.reset_mock()
    try:
        coordinator.fetch_committed_offsets(partitions)
    except AssertionError:
        pass
    else:
        assert False, 'Exception not raised when expected'
    assert coordinator._client.poll.call_count == 1

    coordinator._client.poll.reset_mock()
    coordinator._send_offset_fetch_request.side_effect = [
        Future().failure(Errors.RequestTimedOutError),
        Future().success('fizzbuzz')]

    ret = coordinator.fetch_committed_offsets(partitions)
    assert ret == 'fizzbuzz'
    assert coordinator._client.poll.call_count == 2 # call + retry


def test_close(mocker, coordinator):
    mocker.patch.object(coordinator, '_maybe_auto_commit_offsets_sync')
    mocker.patch.object(coordinator, '_handle_leave_group_response')
    mocker.patch.object(coordinator, 'coordinator_unknown', return_value=False)
    coordinator.coordinator_id = 0
    coordinator._generation = Generation(1, 'foobar', b'')
    coordinator.state = MemberState.STABLE
    cli = coordinator._client
    mocker.patch.object(cli, 'send', return_value=Future().success('foobar'))
    mocker.patch.object(cli, 'poll')

    coordinator.close()
    assert coordinator._maybe_auto_commit_offsets_sync.call_count == 1
    coordinator._handle_leave_group_response.assert_called_with('foobar')

    assert coordinator.generation() is None
    assert coordinator._generation is Generation.NO_GENERATION
    assert coordinator.state is MemberState.UNJOINED
    assert coordinator.rejoin_needed is True


@pytest.fixture
def offsets():
    return {
        TopicPartition('foobar', 0): OffsetAndMetadata(123, b''),
        TopicPartition('foobar', 1): OffsetAndMetadata(234, b''),
    }


def test_commit_offsets_async(mocker, coordinator, offsets):
    mocker.patch.object(coordinator._client, 'poll')
    mocker.patch.object(coordinator, 'coordinator_unknown', return_value=False)
    mocker.patch.object(coordinator, 'ensure_coordinator_ready')
    mocker.patch.object(coordinator, '_send_offset_commit_request',
                        return_value=Future().success('fizzbuzz'))
    coordinator.commit_offsets_async(offsets)
    assert coordinator._send_offset_commit_request.call_count == 1


def test_commit_offsets_sync(mocker, coordinator, offsets):
    mocker.patch.object(coordinator, 'ensure_coordinator_ready')
    mocker.patch.object(coordinator, '_send_offset_commit_request',
                        return_value=Future().success('fizzbuzz'))
    cli = coordinator._client
    mocker.patch.object(cli, 'poll')

    # No offsets, no calls
    assert coordinator.commit_offsets_sync({}) is None
    assert coordinator._send_offset_commit_request.call_count == 0
    assert cli.poll.call_count == 0

    ret = coordinator.commit_offsets_sync(offsets)
    assert coordinator._send_offset_commit_request.call_count == 1
    assert cli.poll.call_count == 1
    assert ret == 'fizzbuzz'

    # Failed future is raised if not retriable
    coordinator._send_offset_commit_request.return_value = Future().failure(AssertionError)
    coordinator._client.poll.reset_mock()
    try:
        coordinator.commit_offsets_sync(offsets)
    except AssertionError:
        pass
    else:
        assert False, 'Exception not raised when expected'
    assert coordinator._client.poll.call_count == 1

    coordinator._client.poll.reset_mock()
    coordinator._send_offset_commit_request.side_effect = [
        Future().failure(Errors.RequestTimedOutError),
        Future().success('fizzbuzz')]

    ret = coordinator.commit_offsets_sync(offsets)
    assert ret == 'fizzbuzz'
    assert coordinator._client.poll.call_count == 2 # call + retry


@pytest.mark.parametrize(
    'api_version,group_id,enable,error,has_auto_commit,commit_offsets,warn,exc', [
        ((0, 8, 0), 'foobar', True, None, False, False, True, False),
        ((0, 8, 1), 'foobar', True, None, True, True, False, False),
        ((0, 8, 2), 'foobar', True, None, True, True, False, False),
        ((0, 9), 'foobar', False, None, False, False, False, False),
        ((0, 9), 'foobar', True, Errors.UnknownMemberIdError(), True, True, True, False),
        ((0, 9), 'foobar', True, Errors.IllegalGenerationError(), True, True, True, False),
        ((0, 9), 'foobar', True, Errors.RebalanceInProgressError(), True, True, True, False),
        ((0, 9), 'foobar', True, Exception(), True, True, False, True),
        ((0, 9), 'foobar', True, None, True, True, False, False),
        ((0, 9), None, True, None, False, False, True, False),
    ])
def test_maybe_auto_commit_offsets_sync(mocker, api_version, group_id, enable,
                                        error, has_auto_commit, commit_offsets,
                                        warn, exc):
    mock_warn = mocker.patch('kafka.coordinator.consumer.log.warning')
    mock_exc = mocker.patch('kafka.coordinator.consumer.log.exception')
    client = KafkaClient(api_version=api_version)
    coordinator = ConsumerCoordinator(client, SubscriptionState(),
                                      Metrics(),
                                      api_version=api_version,
                                      session_timeout_ms=30000,
                                      max_poll_interval_ms=30000,
                                      enable_auto_commit=enable,
                                      group_id=group_id)
    commit_sync = mocker.patch.object(coordinator, 'commit_offsets_sync',
                                      side_effect=error)
    if has_auto_commit:
        assert coordinator.next_auto_commit_deadline is not None
    else:
        assert coordinator.next_auto_commit_deadline is None

    assert coordinator._maybe_auto_commit_offsets_sync() is None

    if has_auto_commit:
        assert coordinator.next_auto_commit_deadline is not None

    assert commit_sync.call_count == (1 if commit_offsets else 0)
    assert mock_warn.call_count == (1 if warn else 0)
    assert mock_exc.call_count == (1 if exc else 0)


@pytest.fixture
def patched_coord(mocker, coordinator):
    coordinator._subscription.subscribe(topics=['foobar'])
    mocker.patch.object(coordinator, 'coordinator_unknown', return_value=False)
    coordinator.coordinator_id = 0
    mocker.patch.object(coordinator, 'coordinator', return_value=0)
    coordinator._generation = Generation(0, 'foobar', b'')
    coordinator.state = MemberState.STABLE
    coordinator.rejoin_needed = False
    mocker.patch.object(coordinator, 'need_rejoin', return_value=False)
    mocker.patch.object(coordinator._client, 'least_loaded_node',
                        return_value=1)
    mocker.patch.object(coordinator._client, 'ready', return_value=True)
    mocker.patch.object(coordinator._client, 'send')
    mocker.patch.object(coordinator, '_heartbeat_thread')
    mocker.spy(coordinator, '_failed_request')
    mocker.spy(coordinator, '_handle_offset_commit_response')
    mocker.spy(coordinator, '_handle_offset_fetch_response')
    return coordinator


def test_send_offset_commit_request_fail(mocker, patched_coord, offsets):
    patched_coord.coordinator_unknown.return_value = True
    patched_coord.coordinator_id = None
    patched_coord.coordinator.return_value = None

    # No offsets
    ret = patched_coord._send_offset_commit_request({})
    assert isinstance(ret, Future)
    assert ret.succeeded()

    # No coordinator
    ret = patched_coord._send_offset_commit_request(offsets)
    assert ret.failed()
    assert isinstance(ret.exception, Errors.GroupCoordinatorNotAvailableError)


@pytest.mark.parametrize('api_version,req_type', [
    ((0, 8, 1), OffsetCommitRequest[0]),
    ((0, 8, 2), OffsetCommitRequest[1]),
    ((0, 9), OffsetCommitRequest[2])])
def test_send_offset_commit_request_versions(patched_coord, offsets,
                                             api_version, req_type):
    expect_node = 0
    patched_coord.config['api_version'] = api_version

    patched_coord._send_offset_commit_request(offsets)
    (node, request), _ = patched_coord._client.send.call_args
    assert node == expect_node, 'Unexpected coordinator node'
    assert isinstance(request, req_type)


def test_send_offset_commit_request_failure(patched_coord, offsets):
    _f = Future()
    patched_coord._client.send.return_value = _f
    future = patched_coord._send_offset_commit_request(offsets)
    (node, request), _ = patched_coord._client.send.call_args
    error = Exception()
    _f.failure(error)
    patched_coord._failed_request.assert_called_with(0, request, future, error)
    assert future.failed()
    assert future.exception is error


def test_send_offset_commit_request_success(mocker, patched_coord, offsets):
    _f = Future()
    patched_coord._client.send.return_value = _f
    future = patched_coord._send_offset_commit_request(offsets)
    (node, request), _ = patched_coord._client.send.call_args
    response = OffsetCommitResponse[0]([('foobar', [(0, 0), (1, 0)])])
    _f.success(response)
    patched_coord._handle_offset_commit_response.assert_called_with(
        offsets, future, mocker.ANY, response)


@pytest.mark.parametrize('response,error,dead', [
    (OffsetCommitResponse[0]([('foobar', [(0, 30), (1, 30)])]),
     Errors.GroupAuthorizationFailedError, False),
    (OffsetCommitResponse[0]([('foobar', [(0, 12), (1, 12)])]),
     Errors.OffsetMetadataTooLargeError, False),
    (OffsetCommitResponse[0]([('foobar', [(0, 28), (1, 28)])]),
     Errors.InvalidCommitOffsetSizeError, False),
    (OffsetCommitResponse[0]([('foobar', [(0, 14), (1, 14)])]),
     Errors.GroupLoadInProgressError, False),
    (OffsetCommitResponse[0]([('foobar', [(0, 15), (1, 15)])]),
     Errors.GroupCoordinatorNotAvailableError, True),
    (OffsetCommitResponse[0]([('foobar', [(0, 16), (1, 16)])]),
     Errors.NotCoordinatorForGroupError, True),
    (OffsetCommitResponse[0]([('foobar', [(0, 7), (1, 7)])]),
     Errors.RequestTimedOutError, True),
    (OffsetCommitResponse[0]([('foobar', [(0, 25), (1, 25)])]),
     Errors.CommitFailedError, False),
    (OffsetCommitResponse[0]([('foobar', [(0, 22), (1, 22)])]),
     Errors.CommitFailedError, False),
    (OffsetCommitResponse[0]([('foobar', [(0, 27), (1, 27)])]),
     Errors.CommitFailedError, False),
    (OffsetCommitResponse[0]([('foobar', [(0, 17), (1, 17)])]),
     Errors.InvalidTopicError, False),
    (OffsetCommitResponse[0]([('foobar', [(0, 29), (1, 29)])]),
     Errors.TopicAuthorizationFailedError, False),
])
def test_handle_offset_commit_response(mocker, patched_coord, offsets,
                                       response, error, dead):
    future = Future()
    patched_coord._handle_offset_commit_response(offsets, future, time.time(),
                                                 response)
    assert isinstance(future.exception, error)
    assert patched_coord.coordinator_id is (None if dead else 0)


@pytest.fixture
def partitions():
    return [TopicPartition('foobar', 0), TopicPartition('foobar', 1)]


def test_send_offset_fetch_request_fail(mocker, patched_coord, partitions):
    patched_coord.coordinator_unknown.return_value = True
    patched_coord.coordinator_id = None
    patched_coord.coordinator.return_value = None

    # No partitions
    ret = patched_coord._send_offset_fetch_request([])
    assert isinstance(ret, Future)
    assert ret.succeeded()
    assert ret.value == {}

    # No coordinator
    ret = patched_coord._send_offset_fetch_request(partitions)
    assert ret.failed()
    assert isinstance(ret.exception, Errors.GroupCoordinatorNotAvailableError)


@pytest.mark.parametrize('api_version,req_type', [
    ((0, 8, 1), OffsetFetchRequest[0]),
    ((0, 8, 2), OffsetFetchRequest[1]),
    ((0, 9), OffsetFetchRequest[1])])
def test_send_offset_fetch_request_versions(patched_coord, partitions,
                                            api_version, req_type):
    # assuming fixture sets coordinator=0, least_loaded_node=1
    expect_node = 0
    patched_coord.config['api_version'] = api_version

    patched_coord._send_offset_fetch_request(partitions)
    (node, request), _ = patched_coord._client.send.call_args
    assert node == expect_node, 'Unexpected coordinator node'
    assert isinstance(request, req_type)


def test_send_offset_fetch_request_failure(patched_coord, partitions):
    _f = Future()
    patched_coord._client.send.return_value = _f
    future = patched_coord._send_offset_fetch_request(partitions)
    (node, request), _ = patched_coord._client.send.call_args
    error = Exception()
    _f.failure(error)
    patched_coord._failed_request.assert_called_with(0, request, future, error)
    assert future.failed()
    assert future.exception is error


def test_send_offset_fetch_request_success(patched_coord, partitions):
    _f = Future()
    patched_coord._client.send.return_value = _f
    future = patched_coord._send_offset_fetch_request(partitions)
    (node, request), _ = patched_coord._client.send.call_args
    response = OffsetFetchResponse[0]([('foobar', [(0, 123, b'', 0), (1, 234, b'', 0)])])
    _f.success(response)
    patched_coord._handle_offset_fetch_response.assert_called_with(
        future, response)


@pytest.mark.parametrize('response,error,dead', [
    (OffsetFetchResponse[0]([('foobar', [(0, 123, b'', 14), (1, 234, b'', 14)])]),
     Errors.GroupLoadInProgressError, False),
    (OffsetFetchResponse[0]([('foobar', [(0, 123, b'', 16), (1, 234, b'', 16)])]),
     Errors.NotCoordinatorForGroupError, True),
    (OffsetFetchResponse[0]([('foobar', [(0, 123, b'', 25), (1, 234, b'', 25)])]),
     Errors.UnknownMemberIdError, False),
    (OffsetFetchResponse[0]([('foobar', [(0, 123, b'', 22), (1, 234, b'', 22)])]),
     Errors.IllegalGenerationError, False),
    (OffsetFetchResponse[0]([('foobar', [(0, 123, b'', 29), (1, 234, b'', 29)])]),
     Errors.TopicAuthorizationFailedError, False),
    (OffsetFetchResponse[0]([('foobar', [(0, 123, b'', 0), (1, 234, b'', 0)])]),
     None, False),
])
def test_handle_offset_fetch_response(patched_coord, offsets,
                                      response, error, dead):
    future = Future()
    patched_coord._handle_offset_fetch_response(future, response)
    if error is not None:
        assert isinstance(future.exception, error)
    else:
        assert future.succeeded()
        assert future.value == offsets
    assert patched_coord.coordinator_id is (None if dead else 0)


def test_heartbeat(mocker, patched_coord):
    heartbeat = HeartbeatThread(patched_coord)

    assert not heartbeat.enabled and not heartbeat.closed

    heartbeat.enable()
    assert heartbeat.enabled

    heartbeat.disable()
    assert not heartbeat.enabled

    # heartbeat disables when un-joined
    heartbeat.enable()
    patched_coord.state = MemberState.UNJOINED
    heartbeat._run_once()
    assert not heartbeat.enabled

    heartbeat.enable()
    patched_coord.state = MemberState.STABLE
    mocker.spy(patched_coord, '_send_heartbeat_request')
    mocker.patch.object(patched_coord.heartbeat, 'should_heartbeat', return_value=True)
    heartbeat._run_once()
    assert patched_coord._send_heartbeat_request.call_count == 1

    heartbeat.close()
    assert heartbeat.closed


def test_lookup_coordinator_failure(mocker, coordinator):

    mocker.patch.object(coordinator, '_send_group_coordinator_request',
                        return_value=Future().failure(Exception('foobar')))
    future = coordinator.lookup_coordinator()
    assert future.failed()


def test_ensure_active_group(mocker, coordinator):
    coordinator._subscription.subscribe(topics=['foobar'])
    mocker.patch.object(coordinator, 'coordinator_unknown', return_value=False)
    mocker.patch.object(coordinator, '_send_join_group_request', return_value=Future().success(True))
    mocker.patch.object(coordinator, 'need_rejoin', side_effect=[True, False])
    mocker.patch.object(coordinator, '_on_join_complete')
    mocker.patch.object(coordinator, '_heartbeat_thread')

    coordinator.ensure_active_group()

    coordinator._send_join_group_request.assert_called_once_with()