1  import asyncio
   2  import contextvars
   3  import time
   4  
   5  from .compatibility import guarantee_single_callable
   6  from .timeout import timeout as async_timeout
   7  
   8  
   9  class ApplicationCommunicator:
  10      """
  11      Runs an ASGI application in a test mode, allowing sending of
  12      messages to it and retrieval of messages it sends.
  13      """
  14  
  15      def __init__(self, application, scope):
  16          self.application = guarantee_single_callable(application)
  17          self.scope = scope
  18          self.input_queue = asyncio.Queue()
  19          self.output_queue = asyncio.Queue()
  20          # Clear context - this ensures that context vars set in the testing scope
  21          # are not "leaked" into the application which would normally begin with
  22          # an empty context. In Python >= 3.11 this could also be written as:
  23          # asyncio.create_task(..., context=contextvars.Context())
  24          self.future = contextvars.Context().run(
  25              asyncio.create_task,
  26              self.application(scope, self.input_queue.get, self.output_queue.put),
  27          )
  28  
  29      async def wait(self, timeout=1):
  30          """
  31          Waits for the application to stop itself and returns any exceptions.
  32          """
  33          try:
  34              async with async_timeout(timeout):
  35                  try:
  36                      await self.future
  37                      self.future.result()
  38                  except asyncio.CancelledError:
  39                      pass
  40          finally:
  41              if not self.future.done():
  42                  self.future.cancel()
  43                  try:
  44                      await self.future
  45                  except asyncio.CancelledError:
  46                      pass
  47  
  48      def stop(self, exceptions=True):
  49          if not self.future.done():
  50              self.future.cancel()
  51          elif exceptions:
  52              # Give a chance to raise any exceptions
  53              self.future.result()
  54  
  55      def __del__(self):
  56          # Clean up on deletion
  57          try:
  58              self.stop(exceptions=False)
  59          except RuntimeError:
  60              # Event loop already stopped
  61              pass
  62  
  63      async def send_input(self, message):
  64          """
  65          Sends a single message to the application
  66          """
  67          # Give it the message
  68          await self.input_queue.put(message)
  69  
  70      async def receive_output(self, timeout=1):
  71          """
  72          Receives a single message from the application, with optional timeout.
  73          """
  74          # Make sure there's not an exception to raise from the task
  75          if self.future.done():
  76              self.future.result()
  77          # Wait and receive the message
  78          try:
  79              async with async_timeout(timeout):
  80                  return await self.output_queue.get()
  81          except asyncio.TimeoutError as e:
  82              # See if we have another error to raise inside
  83              if self.future.done():
  84                  self.future.result()
  85              else:
  86                  self.future.cancel()
  87                  try:
  88                      await self.future
  89                  except asyncio.CancelledError:
  90                      pass
  91              raise e
  92  
  93      async def receive_nothing(self, timeout=0.1, interval=0.01):
  94          """
  95          Checks that there is no message to receive in the given time.
  96          """
  97          # `interval` has precedence over `timeout`
  98          start = time.monotonic()
  99          while time.monotonic() - start < timeout:
 100              if not self.output_queue.empty():
 101                  return False
 102              await asyncio.sleep(interval)
 103          return self.output_queue.empty()
 104