def test_scale_after_rolling_update(
docs, replicas, scale_to, expected_before_scale, expected_after_scale
):
flow = Flow().add(
name='executor1',
uses=DummyMarkExecutor,
replicas=replicas,
)
with flow:
ret1 = flow.search(docs, return_results=True, request_size=1)
flow.rolling_update('executor1', None)
flow.scale('executor1', replicas=scale_to)
ret2 = flow.search(docs, return_results=True, request_size=1)
replica_ids = set()
for r in ret1:
for replica_id in r.docs.get_attributes('tags__replica'):
replica_ids.add(replica_id)
assert replica_ids == expected_before_scale
replica_ids = set()
for r in ret2:
for replica_id in r.docs.get_attributes('tags__replica'):
replica_ids.add(replica_id)
assert replica_ids == expected_after_scale
def test_scale_after_rolling_update(docs, replicas, scale_to):
flow = Flow(port_expose=exposed_port).add(
name='executor1',
uses=DummyMarkExecutor,
replicas=replicas,
)
with flow:
ret1 = Client(port=exposed_port).search(docs,
return_results=True,
request_size=1)
flow.rolling_update('executor1', None)
flow.scale('executor1', replicas=scale_to)
ret2 = Client(port=exposed_port).search(docs,
return_results=True,
request_size=1)
replicas_before = set()
for r in ret1:
for replica in r.docs[:, 'tags__replica']:
replicas_before.add(replica)
assert len(replicas_before) == replicas
replicas_after = set()
for r in ret2:
for replica in r.docs[:, 'tags__replica']:
replicas_after.add(replica)
assert len(replicas_after) == scale_to
