def test_get_async_invocation_wrapper_no_extension(self): """The async wrapper will wrap an asynchronous function with a coroutine interface. When there is no extension, it should only invoke the customer's function. """ # Create a mocked customer_function with async wrapper result = aio_compat.run( self._instance.get_async_invocation_wrapper( self._mock_context, self._mock_function_main_async, self._mock_arguments)) # Ensure the return value matches the function method self.assertEqual(result, 'request_ok')
def test_get_async_invocation_wrapper_with_func_extension(self): """The async wrapper will wrap an asynchronous function with a coroutine interface. When there is registered extension, it should execute the extension as well. """ # Register a function extension FuncExtClass = self._generate_new_func_extension_class( self._sdk.FuncExtensionBase, self._mock_func_name) _func_ext_instance = FuncExtClass() # Create a mocked customer_function with async wrapper result = aio_compat.run( self._instance.get_async_invocation_wrapper( self._mock_context, self._mock_function_main_async, self._mock_arguments)) # Ensure the extension is executed self.assertTrue(_func_ext_instance._pre_invocation_executed) # Ensure the customer's function is executed self.assertEqual(result, 'request_ok')
def test_get_invocation_async_disabled_with_flag(self): """The async wrapper will only execute customer's function. This should not execute the extension. """ # Turn off feature flag os.environ[PYTHON_ENABLE_WORKER_EXTENSIONS] = 'false' # Register a function extension FuncExtClass = self._generate_new_func_extension_class( self._sdk.FuncExtensionBase, self._mock_func_name) _func_ext_instance = FuncExtClass() # Create a mocked customer_function with async wrapper result = aio_compat.run( self._instance.get_async_invocation_wrapper( self._mock_context, self._mock_function_main_async, self._mock_arguments)) # The extension SHOULD NOT be executed self.assertFalse(_func_ext_instance._pre_invocation_executed) # Ensure the customer's function is executed self.assertEqual(result, 'request_ok')
def wrapper(*args, **kwargs): return aio_compat.run(func(*args, **kwargs))