sientia-dataops-opc-ingestor/e2e/test_slot_leasing.py

"""E2E scenarios — 1. Slot leasing and load balancing (Redis). See ``e2e/scenarios.md``."""

import asyncio

import pytest

from e2e.helpers import (
    load_scenario_input,
    notifications,
    readings,
    redis_json,
    stays_false,
    wait_until,
)

pytestmark = [pytest.mark.integration, pytest.mark.asyncio]


async def _two_slots(opc_servers, put_slot):
    """Slot 1 reads plant-a, slot 2 reads plant-b."""
    plant_a = await opc_servers('plant-a', {'Temperature': 20.0})
    plant_b = await opc_servers('plant-b', {'Flow': 3.0})
    both = load_scenario_input('slot_two_servers.json', {'plant-a': plant_a, 'plant-b': plant_b})
    put_slot(1, {'plant-a': both['plant-a']})
    put_slot(2, {'plant-b': both['plant-b']})
    return plant_a, plant_b


# 1.1 Acquisition ----------------------------------------------------------------------


async def test_scenario_1_1_1_single_pod_leases_the_slot_and_heartbeats(
    opc_servers, pods, put_slot, redis_client, mongo_db
):
    server = await opc_servers('plant-a', {'Temperature': 20.0})
    put_slot(1, load_scenario_input('slot_single_server.json', {'plant-a': server}))

    pod = await pods('pod-a')
    await pod.tick()

    assert pod.slots == {'1'}
    assert redis_json(redis_client, 'lease:opc_tags:1') == 'pod-a'
    assert 0 < redis_client.ttl('lease:opc_tags:1') <= 10
    assert redis_json(redis_client, 'heartbeat:ingestor:pod-a') == 1
    assert 0 < redis_client.ttl('heartbeat:ingestor:pod-a') <= 20
    await wait_until(lambda: readings(mongo_db, 'raw_plant_a'))


async def test_scenario_1_1_2_idle_pod_takes_a_slot_created_after_startup(
    opc_servers, pods, put_slot, redis_client, mongo_db
):
    server = await opc_servers('plant-a', {'Temperature': 20.0})
    pod = await pods('pod-a')
    await pod.tick()
    assert pod.slots == set()

    put_slot(1, load_scenario_input('slot_single_server.json', {'plant-a': server}))
    await pod.tick()

    assert pod.slots == {'1'}
    assert redis_json(redis_client, 'lease:opc_tags:1') == 'pod-a'
    await wait_until(lambda: readings(mongo_db, 'raw_plant_a'))


# 1.2 Balancing and failover ------------------------------------------------------------


async def test_scenario_1_2_1_two_pods_split_two_slots(
    opc_servers, pods, put_slot, redis_client, mongo_db
):
    plant_a, plant_b = await _two_slots(opc_servers, put_slot)

    pod_a = await pods('pod-a')
    await pod_a.tick()
    assert pod_a.slots == {'1', '2'}  # alone: takes every slot

    pod_b = await pods('pod-b')
    await pod_b.tick()
    assert pod_b.slots == set()  # every slot is leased
    await pod_a.tick()  # pod-a now sees a pod without slot and gives one back
    assert pod_a.slots == {'1'}
    await pod_b.tick()
    assert pod_b.slots == {'2'}

    assert redis_json(redis_client, 'lease:opc_tags:1') == 'pod-a'
    assert redis_json(redis_client, 'lease:opc_tags:2') == 'pod-b'
    assert set(pod_a.manager.opc_managers) == {'plant-a'}
    assert set(pod_b.manager.opc_managers) == {'plant-b'}
    await plant_a.set('Temperature', 25.0)
    await plant_b.set('Flow', 4.5)
    await wait_until(lambda: readings(mongo_db, 'raw_plant_a', value=25.0))
    await wait_until(lambda: readings(mongo_db, 'raw_plant_b', value=4.5))


async def test_scenario_1_2_2_extra_pod_stays_idle(opc_servers, pods, put_slot, redis_client):
    await _two_slots(opc_servers, put_slot)
    pod_a, pod_b = await pods('pod-a'), await pods('pod-b')
    await pod_a.tick()
    await pod_b.tick()
    await pod_a.tick()
    await pod_b.tick()

    pod_c = await pods('pod-c')
    for _ in range(2):
        for pod in (pod_a, pod_b, pod_c):
            await pod.tick()

    assert pod_c.slots == set()
    assert pod_a.slots | pod_b.slots == {'1', '2'}
    assert pod_a.slots.isdisjoint(pod_b.slots)
    assert {
        redis_json(redis_client, 'lease:opc_tags:1'),
        redis_json(redis_client, 'lease:opc_tags:2'),
    } == {
        'pod-a',
        'pod-b',
    }


async def test_scenario_1_2_3_failover_after_lease_expiry(
    opc_servers, pods, put_slot, redis_client, mongo_db
):
    server = await opc_servers('plant-a', {'Temperature': 20.0})
    put_slot(1, load_scenario_input('slot_single_server.json', {'plant-a': server}))
    pod_a = await pods('pod-a', LEASE_TTL='2', HEARTBEAT_TTL='2')
    await pod_a.tick()
    pod_b = await pods('pod-b')
    await pod_b.tick()
    assert (pod_a.slots, pod_b.slots) == ({'1'}, set())

    await pod_a.stop()  # crash-like: the lease is not released, it expires
    assert redis_json(redis_client, 'lease:opc_tags:1') == 'pod-a'
    await wait_until(lambda: redis_json(redis_client, 'lease:opc_tags:1') is None, timeout=5)
    await pod_b.tick()

    assert pod_b.slots == {'1'}
    assert redis_json(redis_client, 'lease:opc_tags:1') == 'pod-b'
    await server.set('Temperature', 30.0)
    await wait_until(lambda: readings(mongo_db, 'raw_plant_a', value=30.0))


# 1.3 Slot changes and invalid slots ----------------------------------------------------


async def test_scenario_1_3_1_deleted_slot_stops_the_reads(
    opc_servers, pods, put_slot, redis_client, mongo_db
):
    server = await opc_servers('plant-a', {'Temperature': 20.0})
    put_slot(1, load_scenario_input('slot_single_server.json', {'plant-a': server}))
    pod = await pods('pod-a')
    await pod.tick()
    await wait_until(lambda: readings(mongo_db, 'raw_plant_a'))

    redis_client.delete('slot:opc_tags:1')
    await pod.tick()

    assert pod.slots == set()
    assert pod.manager.opc_managers == {}
    await server.set('Temperature', 99.0)
    assert await stays_false(lambda: readings(mongo_db, 'raw_plant_a', value=99.0))


async def test_scenario_1_3_2_invalid_slot_is_notified_and_not_read(
    opc_servers, pods, put_slot, mongo_db
):
    server = await opc_servers('plant-a', {'Temperature': 20.0})
    put_slot(1, load_scenario_input('slot_v1_invalid.json', {'plant-a': server}))

    pod = await pods('pod-a')
    await pod.tick()

    assert pod.slots == set()
    assert notifications(mongo_db, 'OPC_SLOT_CONFIG_ERROR_1')
    assert await stays_false(lambda: readings(mongo_db, 'raw_plant_a'))


async def test_scenario_1_3_3_fixed_slot_is_read_once_its_lease_expires(
    opc_servers, pods, put_slot, mongo_db
):
    server = await opc_servers('plant-a', {'Temperature': 20.0})
    put_slot(1, load_scenario_input('slot_v1_invalid.json', {'plant-a': server}))
    pod = await pods('pod-a', LEASE_TTL='2')
    await pod.tick()
    assert pod.slots == set()

    put_slot(1, load_scenario_input('slot_single_server.json', {'plant-a': server}))
    await pod.tick()
    assert pod.slots == set()  # still leased by the failed acquisition

    await asyncio.sleep(2.5)
    await pod.tick()
    assert pod.slots == {'1'}
    await wait_until(lambda: readings(mongo_db, 'raw_plant_a'))