from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.plugins import (
    MessageSink,
    MessageSource,
    expect,
)


class SendPing(MessageSource, MessageSink):
    async def create_sink(self, loop):
        self._loop = loop

    async def create_source(self, loop, sink):
        self._loop = loop
        self._sink = sink

    @expect(MessageType.ServerConnected, MessageType.ServerReconnected)
    async def send_ping(self, _):
        await self._sink.process_message(MessageType.Ping())