import unittest

from ipython_shell.transport import JupyterTransport, KernelMessage


class TransportTests(unittest.TestCase):
    def test_kernel_message_keeps_content_and_buffers(self):
        message = KernelMessage(
            msg_type="display_data",
            parent_id="call-1",
            content={"data": {"text/plain": "42"}},
            buffers=[b"binary"],
        )

        self.assertEqual(message.parent_id, "call-1")
        self.assertEqual(message.content["data"], {"text/plain": "42"})
        self.assertEqual(message.buffers, [b"binary"])


class AsyncTransportTests(unittest.IsolatedAsyncioTestCase):
    async def test_transport_executes_and_returns_execute_reply(self):
        transport = JupyterTransport()
        await transport.start()
        try:
            call_id = await transport.execute("2 + 2")
            messages = [
                message async for message in transport.messages_for(call_id)
            ]
        finally:
            await transport.shutdown()

        self.assertTrue(
            any(message.msg_type == "execute_result" for message in messages)
        )
        self.assertEqual(messages[-1].msg_type, "execute_reply")
        self.assertEqual(messages[-1].content["status"], "ok")


if __name__ == "__main__":
    unittest.main()