assigner_impl_test.py

import asyncio
from unittest.mock import call
from collections import defaultdict
from typing import Dict, Set

from asynctest.mock import MagicMock, CoroutineMock
import pytest

from google.cloud.pubsublite.internal.wire.assigner import Assigner
from google.cloud.pubsublite.internal.wire.assigner_impl import AssignerImpl
from google.cloud.pubsublite.internal.wire.connection import Connection, ConnectionFactory
from google.api_core.exceptions import InternalServerError

from google.cloud.pubsublite.partition import Partition
from google.cloud.pubsublite_v1.types.subscriber import PartitionAssignmentRequest, InitialPartitionAssignmentRequest, \
  PartitionAssignment, PartitionAssignmentAck
from google.cloud.pubsublite.testing.test_utils import make_queue_waiter
from google.cloud.pubsublite.internal.wire.retrying_connection import _MIN_BACKOFF_SECS

# All test coroutines will be treated as marked.
pytestmark = pytest.mark.asyncio


@pytest.fixture()
def default_connection():
  conn = MagicMock(spec=Connection[PartitionAssignmentRequest, PartitionAssignment])
  conn.__aenter__.return_value = conn
  return conn


@pytest.fixture()
def connection_factory(default_connection):
  factory = MagicMock(spec=ConnectionFactory[PartitionAssignmentRequest, PartitionAssignment])
  factory.new.return_value = default_connection
  return factory


@pytest.fixture()
def initial_request():
  return PartitionAssignmentRequest(initial=InitialPartitionAssignmentRequest(subscription="mysub"))


class QueuePair:
  called: asyncio.Queue
  results: asyncio.Queue

  def __init__(self):
    self.called = asyncio.Queue()
    self.results = asyncio.Queue()


@pytest.fixture
def sleep_queues() -> Dict[float, QueuePair]:
  return defaultdict(QueuePair)


@pytest.fixture
def asyncio_sleep(monkeypatch, sleep_queues):
  """Requests.get() mocked to return {'mock_key':'mock_response'}."""
  mock = CoroutineMock()
  monkeypatch.setattr(asyncio, "sleep", mock)

  async def sleeper(delay: float):
    await make_queue_waiter(sleep_queues[delay].called, sleep_queues[delay].results)(delay)

  mock.side_effect = sleeper
  return mock


@pytest.fixture()
def assigner(connection_factory, initial_request):
  return AssignerImpl(initial_request.initial, connection_factory)


def as_response(partitions: Set[Partition]):
  req = PartitionAssignment()
  req.partitions = [partition.value for partition in partitions]
  return req


def ack_request():
  return PartitionAssignmentRequest(ack=PartitionAssignmentAck())


async def test_basic_assign(
    assigner: Assigner, default_connection, initial_request):
  write_called_queue = asyncio.Queue()
  write_result_queue = asyncio.Queue()
  default_connection.write.side_effect = make_queue_waiter(write_called_queue, write_result_queue)
  read_called_queue = asyncio.Queue()
  read_result_queue = asyncio.Queue()
  default_connection.read.side_effect = make_queue_waiter(read_called_queue, read_result_queue)
  write_result_queue.put_nowait(None)
  async with assigner:
    # Set up connection
    await write_called_queue.get()
    await read_called_queue.get()
    default_connection.write.assert_has_calls([call(initial_request)])

    # Wait for the first assignment
    assign_fut1 = asyncio.ensure_future(assigner.get_assignment())
    assert not assign_fut1.done()

    partitions = {Partition(2), Partition(7)}

    # Send the first assignment.
    await read_result_queue.put(as_response(partitions=partitions))
    assert (await assign_fut1) == partitions

    # Get the next assignment: should send an ack on the stream
    assign_fut2 = asyncio.ensure_future(assigner.get_assignment())
    await write_called_queue.get()
    await write_result_queue.put(None)
    default_connection.write.assert_has_calls([call(initial_request), call(ack_request())])

    partitions = {Partition(5)}

    # Send the second assignment.
    await read_called_queue.get()
    await read_result_queue.put(as_response(partitions=partitions))
    assert (await assign_fut2) == partitions


async def test_restart(
    assigner: Assigner, default_connection, connection_factory, initial_request, asyncio_sleep, sleep_queues):
  write_called_queue = asyncio.Queue()
  write_result_queue = asyncio.Queue()
  default_connection.write.side_effect = make_queue_waiter(write_called_queue, write_result_queue)
  read_called_queue = asyncio.Queue()
  read_result_queue = asyncio.Queue()
  default_connection.read.side_effect = make_queue_waiter(read_called_queue, read_result_queue)
  write_result_queue.put_nowait(None)
  async with assigner:
    # Set up connection
    await write_called_queue.get()
    await read_called_queue.get()
    default_connection.write.assert_has_calls([call(initial_request)])

    # Wait for the first assignment
    assign_fut1 = asyncio.ensure_future(assigner.get_assignment())
    assert not assign_fut1.done()

    partitions = {Partition(2), Partition(7)}

    # Send the first assignment.
    await read_result_queue.put(as_response(partitions=partitions))
    await read_called_queue.get()
    assert (await assign_fut1) == partitions

    # Get the next assignment: should attempt to send an ack on the stream
    assign_fut2 = asyncio.ensure_future(assigner.get_assignment())
    await write_called_queue.get()
    default_connection.write.assert_has_calls([call(initial_request), call(ack_request())])

    # Set up the next connection
    conn2 = MagicMock(spec=Connection[PartitionAssignmentRequest, PartitionAssignment])
    conn2.__aenter__.return_value = conn2
    connection_factory.new.return_value = conn2
    write_called_queue_2 = asyncio.Queue()
    write_result_queue_2 = asyncio.Queue()
    conn2.write.side_effect = make_queue_waiter(write_called_queue_2, write_result_queue_2)
    read_called_queue_2 = asyncio.Queue()
    read_result_queue_2 = asyncio.Queue()
    conn2.read.side_effect = make_queue_waiter(read_called_queue_2, read_result_queue_2)

    # Fail the connection by failing the write call.
    await write_result_queue.put(InternalServerError("failed"))
    await sleep_queues[_MIN_BACKOFF_SECS].called.get()
    await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)

    # Reinitialize
    await write_called_queue_2.get()
    write_result_queue_2.put_nowait(None)
    conn2.write.assert_has_calls([call(initial_request)])

    partitions = {Partition(5)}

    # Send the second assignment on the new connection.
    await read_called_queue_2.get()
    await read_result_queue_2.put(as_response(partitions=partitions))
    assert (await assign_fut2) == partitions
    # No ack call ever made.
    conn2.write.assert_has_calls([call(initial_request)])