code: def test_generator_returns(
      ray_start_regular, use_actors, store_in_plasma
  ):
      remote_generator_fn = None
      if use_actors:
  
          @ray.remote
          class Generator:
              def __init__(self):
                  pass
  
              def generator(self, num_returns, store_in_plasma):
                  for i in range(num_returns):
                      if store_in_plasma:
                          yield np.ones(1_000_000, dtype=np.int8) * i
                      else:
                          yield [i]
  
          g = Generator.remote()
          remote_generator_fn = g.generator
      else:
  
          @ray.remote(max_retries=0)
          def generator(num_returns, store_in_plasma):
              for i in range(num_returns):
                  if store_in_plasma:
                      yield np.ones(1_000_000, dtype=np.int8) * i
                  else:
                      yield [i]
  
          remote_generator_fn = generator
  
      # Check cases when num_returns does not match the number of values returned
      # by the generator.
      num_returns = 3
  
      try:
          ray.get(
              remote_generator_fn.options(
                  num_returns=num_returns
              ).remote(num_returns - 1, store_in_plasma)
          )
          assert False
      except ray.exceptions.RayTaskError as e:
          assert isinstance(e.as_instanceof_cause(), ValueError)
  
      # TODO(swang): When generators return more values than expected, we log an
      # error but the exception is not thrown to the application.
      # https://github.com/ray-project/ray/issues/28689.
      ray.get(
          remote_generator_fn.options(num_returns=num_returns).remote(
              num_returns + 1, store_in_plasma
          )
      )
  
      # Check return values.
      [
          x[0]
          for x in ray.get(
              remote_generator_fn.options(
                  num_returns=num_returns
              ).remote(num_returns, store_in_plasma)
          )
      ] == list(range(num_returns))
      # Works for num_returns=1 if generator returns a single value.
      assert (
          ray.get(
              remote_generator_fn.options(num_returns=1).remote(
                  1, store_in_plasma
              )
          )[0]
          == 0
      )
  
imports: import pytest
  import numpy
  import sys
  import time
  import gc
  from unittest.mock import Mock
  import ray
  from ray.util.client.ray_client_helpers import ray_start_client_server_for_address
  from ray._private.client_mode_hook import enable_client_mode
  from ray.tests.conftest import call_ray_start_context
label_desc: Use image node_id 3 for calling a global function with example usage: test_generator_returns(ray_start_regular, use_actors, store_in_plasma) without return types