async def test_smoke(self):
ni = 0
async with MockAdminAPI() as client:
nodes_config_resp = await client.getNodesConfig(
NodesFilter(node=NodeID(node_index=ni))
)
nodes_state_resp = await client.getNodesState(
NodesStateRequest(filter=NodesFilter(node=NodeID(node_index=ni)))
)
maintenances_resp = await client.getMaintenances(MaintenancesFilter())
nc = nodes_config_resp.nodes[0]
ns = nodes_state_resp.states[0]
mnt_ids = set()
for mnt in maintenances_resp.maintenances:
for s in mnt.shards:
if s.node.node_index == ni:
mnt_ids.add(mnt.group_id)
for n in mnt.sequencer_nodes:
if n.node_index == ni:
mnt_ids.add(mnt.group_id)
mnts = tuple(
sorted(
(
mnt
for mnt in maintenances_resp.maintenances
if mnt.group_id in mnt_ids
),
key=operator.attrgetter("group_id"),
)
)
nv = NodeView(node_config=nc, node_state=ns, maintenances=mnts)
self._validate(nv, nc, ns, mnts)
async def test_get_sequencer_last_updated_at(self):
async with MockAdminAPI() as client:
cv = await get_cluster_view(client)
node_id = cv.get_node_view(node_index=0).node_id
await apply_maintenance(client=client, sequencer_nodes=[node_id])
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertIsNone(mv.get_sequencer_last_updated_at(node_id))
ts = datetime.now()
client._set_sequencer_maintenance_progress(
node_id=node_id,
maintenance_progress=SequencerMaintenanceProgress(
status=MaintenanceStatus.STARTED,
target_state=SequencingState.UNKNOWN,
created_at=datetime.now(),
last_updated_at=datetime.now(),
associated_group_ids=["johnsnow"],
).to_thrift(),
)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertTrue(
mv.get_sequencer_last_updated_at(node_id) -
ts < timedelta(seconds=2))
async def test_smoke(self):
async with MockAdminAPI() as client:
cv = await get_cluster_view(client)
await apply_maintenance(
client=client,
shards=[
ShardID(node=cv.get_node_view_by_node_index(0).node_id,
shard_index=1)
],
sequencer_nodes=[cv.get_node_view_by_node_index(0).node_id],
)
await apply_maintenance(
client=client,
node_ids=[cv.get_node_id(node_index=1)],
user="******",
reason="whatever",
)
(cv, nc_resp, ns_resp, mnts_resp) = await asyncio.gather(
get_cluster_view(client),
client.getNodesConfig(NodesFilter()),
client.getNodesState(NodesStateRequest()),
client.getMaintenances(MaintenancesFilter()),
)
self._validate(cv, nc_resp.nodes, ns_resp.states,
tuple(mnts_resp.maintenances))
async def test_get_shard_last_updated_at(self):
async with MockAdminAPI() as client:
cv = await get_cluster_view(client)
shard = ShardID(node=cv.get_node_view_by_node_index(0).node_id,
shard_index=1)
await apply_maintenance(
client=client,
shards=[shard],
shard_target_state=ShardOperationalState.DRAINED,
)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertIsNone(mv.get_shard_last_updated_at(shard))
ts = datetime.now()
client._set_shard_maintenance_progress(
shard,
ShardMaintenanceProgress(
status=MaintenanceStatus.STARTED,
target_states=[ShardOperationalState.MAY_DISAPPEAR],
created_at=ts,
last_updated_at=datetime.now(),
associated_group_ids=["johnsnow"],
).to_thrift(),
)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertTrue(
mv.get_shard_last_updated_at(shard) -
ts < timedelta(seconds=2))
async def test_sequencer_only(self):
async with MockAdminAPI() as client:
cv = await get_cluster_view(client)
await apply_maintenance(
client=client,
sequencer_nodes=[cv.get_node_view_by_node_index(0).node_id],
)
cv = await get_cluster_view(client)
self.validate(
maintenance_view=list(cv.get_all_maintenance_views())[0],
maintenance=list(cv.get_all_maintenances())[0],
node_index_to_node_view={0: cv.get_node_view(node_index=0)},
)
async def test_smoke(self):
ni = 0
async with MockAdminAPI() as client:
cv = await get_cluster_view(client)
maintenances_resp = await apply_maintenance(
client=client,
shards=[
ShardID(
node=cv.get_node_view_by_node_index(0).node_id, shard_index=1
)
],
sequencer_nodes=[cv.get_node_view_by_node_index(0).node_id],
)
(
nodes_config_resp,
nodes_state_resp,
maintenances_resp,
) = await asyncio.gather(
client.getNodesConfig(NodesFilter(node=NodeID(node_index=ni))),
client.getNodesState(
NodesStateRequest(filter=NodesFilter(node=NodeID(node_index=ni)))
),
client.getMaintenances(MaintenancesFilter()),
)
nc = [n for n in nodes_config_resp.nodes if n.node_index == ni][0]
ns = [n for n in nodes_state_resp.states if n.node_index == ni][0]
mnt_ids = set()
for mnt in maintenances_resp.maintenances:
for s in mnt.shards:
if s.node.node_index == ni:
mnt_ids.add(mnt.group_id)
for n in mnt.sequencer_nodes:
if n.node_index == ni:
mnt_ids.add(mnt.group_id)
mnts = tuple(
sorted(
(
mnt
for mnt in maintenances_resp.maintenances
if mnt.group_id in mnt_ids
),
key=operator.attrgetter("group_id"),
)
)
nv = NodeView(node_config=nc, node_state=ns, maintenances=mnts)
self._validate(nv, nc, ns, mnts)
async def test_no_timestamps(self):
async with MockAdminAPI() as client:
cv = await get_cluster_view(client)
await apply_maintenance(
client=client,
sequencer_nodes=[cv.get_node_view_by_node_index(0).node_id],
ttl=timedelta(days=2),
)
cv = await get_cluster_view(client)
mnt = list(cv.get_all_maintenances())[0]
mnt = mnt(created_on=None, expires_on=None)
mv = MaintenanceView(
maintenance=mnt,
node_index_to_node_view={0: cv.get_node_view(node_index=0)})
self.assertIsNone(mv.created_on)
self.assertIsNone(mv.expires_on)
self.assertIsNone(mv.expires_in)
async def test_mismatch(self):
async with MockAdminAPI() as client:
(
nodes_config_resp,
nodes_state_resp,
maintenances_resp,
) = await asyncio.gather(
client.getNodesConfig(NodesFilter(node=NodeID(node_index=0))),
client.getNodesState(
NodesStateRequest(filter=NodesFilter(node=NodeID(node_index=1)))
),
client.getMaintenances(MaintenancesFilter()),
)
with self.assertRaises(ValueError):
NodeView(
node_config=nodes_config_resp.nodes[0],
node_state=nodes_state_resp.states[0],
maintenances=maintenances_resp.maintenances,
)
async def test_node_is_not_a_sequencer(self):
async with MockAdminAPI(disaggregated=True) as client:
cv = await get_cluster_view(client)
shard = ShardID(node=cv.get_node_view_by_node_index(0).node_id,
shard_index=1)
await apply_maintenance(
client=client,
shards=[shard],
shard_target_state=ShardOperationalState.DRAINED,
)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
node_id = cv.get_node_view_by_node_index(0).node_id
with self.assertRaises(NodeIsNotASequencerError):
mv.get_sequencer_maintenance_status(node_id)
with self.assertRaises(NodeIsNotASequencerError):
mv.get_sequencer_last_updated_at(node_id)
async def test_node_is_sequencer_only(self):
async with MockAdminAPI(disaggregated=True, num_sequencer_nodes=1) as client:
cv = await get_cluster_view(client)
# A sequencers in MockAdminAPI start from the node_index >=
# num_storage_nodes
node_id = cv.get_node_view_by_node_index(client.num_storage_nodes).node_id
# We are applying a shard maintenance even that this node doesn't
# have shards (sequencer-only)
shard = ShardID(node=node_id, shard_index=1)
await apply_maintenance(
client=client,
shards=[shard],
# sequencer_nodes=[node_id],
shard_target_state=ShardOperationalState.DRAINED,
)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(mv.get_shard_state(shard), None)
self.assertEqual(
mv.get_shard_maintenance_status(shard), MaintenanceStatus.COMPLETED
)
self.assertEqual(mv.get_shard_last_updated_at(shard), None)
async def test_sequencer_maintenance_status(self):
async with MockAdminAPI() as client:
cv = await get_cluster_view(client)
node_id = cv.get_node_view(node_index=0).node_id
await apply_maintenance(client=client, sequencer_nodes=[node_id])
# Just started
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(
mv.get_sequencer_maintenance_status(node_id),
MaintenanceStatus.NOT_STARTED,
)
self.assertEqual(mv.num_sequencers_total, 1)
self.assertEqual(mv.num_sequencers_done, 0)
self.assertFalse(mv.are_all_sequencers_done)
self.assertFalse(mv.is_everything_done)
self.assertFalse(mv.is_blocked)
self.assertFalse(mv.is_completed)
self.assertTrue(mv.is_in_progress)
self.assertEqual(mv.overall_status,
MaintenanceOverallStatus.IN_PROGRESS)
# In progress
client._set_sequencer_maintenance_progress(
node_id=node_id,
maintenance_progress=SequencerMaintenanceProgress(
status=MaintenanceStatus.STARTED,
target_state=SequencingState.UNKNOWN,
created_at=datetime.now(),
last_updated_at=datetime.now(),
associated_group_ids=["johnsnow"],
).to_thrift(),
)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(mv.get_sequencer_maintenance_status(node_id),
MaintenanceStatus.STARTED)
self.assertEqual(mv.num_sequencers_done, 0)
self.assertFalse(mv.are_all_sequencers_done)
self.assertFalse(mv.is_everything_done)
self.assertFalse(mv.is_blocked)
self.assertFalse(mv.is_completed)
self.assertTrue(mv.is_in_progress)
self.assertEqual(mv.overall_status,
MaintenanceOverallStatus.IN_PROGRESS)
# Blocked
client._set_sequencer_maintenance_progress(
node_id=node_id,
maintenance_progress=SequencerMaintenanceProgress(
status=MaintenanceStatus.BLOCKED_UNTIL_SAFE,
target_state=SequencingState.UNKNOWN,
created_at=datetime.now(),
last_updated_at=datetime.now(),
associated_group_ids=["johnsnow"],
).to_thrift(),
)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(
mv.get_sequencer_maintenance_status(node_id),
MaintenanceStatus.BLOCKED_UNTIL_SAFE,
)
self.assertEqual(mv.num_sequencers_done, 0)
self.assertFalse(mv.are_all_sequencers_done)
self.assertFalse(mv.is_everything_done)
self.assertTrue(mv.is_blocked)
self.assertFalse(mv.is_completed)
self.assertFalse(mv.is_in_progress)
self.assertEqual(mv.overall_status,
MaintenanceOverallStatus.BLOCKED)
# Done
client._set_sequencing_state(node_id, SequencingState.DISABLED)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(
mv.get_sequencer_maintenance_status(node_id),
MaintenanceStatus.COMPLETED,
)
self.assertEqual(mv.overall_status,
MaintenanceOverallStatus.COMPLETED)
async def test_shard_maintenance_status(self):
## MAY_DISAPPEAR maintenance
async with MockAdminAPI() as client:
cv = await get_cluster_view(client)
shard = ShardID(node=cv.get_node_view_by_node_index(0).node_id,
shard_index=1)
await apply_maintenance(
client=client,
shards=[shard],
shard_target_state=ShardOperationalState.MAY_DISAPPEAR,
)
cv = await get_cluster_view(client)
# Just started
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(mv.get_shard_maintenance_status(shard),
MaintenanceStatus.NOT_STARTED)
self.assertEqual(mv.num_shards_done, 0)
self.assertFalse(mv.are_all_shards_done)
self.assertFalse(mv.is_everything_done)
self.assertFalse(mv.is_blocked)
self.assertFalse(mv.is_completed)
self.assertTrue(mv.is_in_progress)
self.assertFalse(mv.is_internal)
self.assertEqual(mv.overall_status,
MaintenanceOverallStatus.IN_PROGRESS)
# In progress
client._set_shard_maintenance_progress(
shard,
ShardMaintenanceProgress(
status=MaintenanceStatus.STARTED,
target_states=[ShardOperationalState.MAY_DISAPPEAR],
created_at=datetime.now(),
last_updated_at=datetime.now(),
associated_group_ids=["johnsnow"],
).to_thrift(),
)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(mv.get_shard_maintenance_status(shard),
MaintenanceStatus.STARTED)
self.assertEqual(mv.num_shards_done, 0)
self.assertFalse(mv.are_all_shards_done)
self.assertFalse(mv.is_everything_done)
self.assertFalse(mv.is_blocked)
self.assertFalse(mv.is_completed)
self.assertTrue(mv.is_in_progress)
self.assertFalse(mv.is_internal)
self.assertEqual(mv.overall_status,
MaintenanceOverallStatus.IN_PROGRESS)
# Blocked
client._set_shard_maintenance_progress(
shard,
ShardMaintenanceProgress(
status=MaintenanceStatus.BLOCKED_UNTIL_SAFE,
target_states=[ShardOperationalState.MAY_DISAPPEAR],
created_at=datetime.now(),
last_updated_at=datetime.now(),
associated_group_ids=["johnsnow"],
).to_thrift(),
)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertTrue(mv.is_blocked)
self.assertFalse(mv.are_all_shards_done)
self.assertFalse(mv.is_everything_done)
self.assertTrue(mv.is_blocked)
self.assertFalse(mv.is_completed)
self.assertFalse(mv.is_in_progress)
self.assertEqual(mv.overall_status,
MaintenanceOverallStatus.BLOCKED)
# Done
for sos in {
ShardOperationalState.DRAINED,
ShardOperationalState.MAY_DISAPPEAR,
ShardOperationalState.MIGRATING_DATA,
ShardOperationalState.PROVISIONING,
}:
client._set_shard_current_operational_state(shard, sos)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(mv.get_shard_maintenance_status(shard),
MaintenanceStatus.COMPLETED)
self.assertEqual(mv.num_shards_done, 1)
self.assertTrue(mv.are_all_shards_done)
self.assertTrue(mv.is_everything_done)
self.assertFalse(mv.is_blocked)
self.assertTrue(mv.is_completed)
self.assertFalse(mv.is_in_progress)
self.assertEqual(mv.overall_status,
MaintenanceOverallStatus.COMPLETED)
## DRAINED maintenance
async with MockAdminAPI() as client:
cv = await get_cluster_view(client)
shard = ShardID(node=cv.get_node_view_by_node_index(0).node_id,
shard_index=1)
await apply_maintenance(
client=client,
shards=[shard],
shard_target_state=ShardOperationalState.DRAINED,
)
cv = await get_cluster_view(client)
# Just started
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(mv.get_shard_maintenance_status(shard),
MaintenanceStatus.NOT_STARTED)
self.assertEqual(mv.num_shards_done, 0)
self.assertFalse(mv.are_all_shards_done)
# May disappear
client._set_shard_current_operational_state(
shard, ShardOperationalState.MAY_DISAPPEAR)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(mv.get_shard_maintenance_status(shard),
MaintenanceStatus.NOT_STARTED)
self.assertEqual(mv.num_shards_done, 0)
self.assertFalse(mv.are_all_shards_done)
# Done
client._set_shard_current_operational_state(
shard, ShardOperationalState.DRAINED)
cv = await get_cluster_view(client)
mv = list(cv.get_all_maintenance_views())[0]
self.assertEqual(mv.get_shard_maintenance_status(shard),
MaintenanceStatus.COMPLETED)
self.assertEqual(mv.num_shards_done, 1)
self.assertTrue(mv.are_all_shards_done)
async def test_disagg(self) -> None:
async with MockAdminAPI(disaggregated=True,
num_storage_nodes=50,
num_sequencer_nodes=50) as client:
return await self.smoke(client)
async def test_normal(self) -> None:
async with MockAdminAPI() as client:
return await self.smoke(client)
