def test_should_use_the_same_function_if_not_gpu_compatible(self):
mock_function = MagicMock(return_value=pd.DataFrame())
expression = FunctionExpression(mock_function, name="test")
input_batch = Batch(frames=pd.DataFrame())
expression.evaluate(input_batch)
mock_function.assert_called()
def test_should_execute_same_function_if_no_gpu(self, context):
context_instance = context.return_value
mock_function = MagicMock(spec=GPUCompatible,
return_value=pd.DataFrame())
context_instance.gpu_device.return_value = NO_GPU
expression = FunctionExpression(mock_function, name="test")
input_batch = Batch(frames=pd.DataFrame())
expression.evaluate(input_batch)
mock_function.assert_called()
def test_function_move_the_device_to_gpu_if_compatible(self, context):
context_instance = context.return_value
mock_function = MagicMock(spec=GPUCompatible)
gpu_mock_function = Mock(return_value=pd.DataFrame())
gpu_device_id = '2'
mock_function.to_device.return_value = gpu_mock_function
context_instance.gpu_device.return_value = gpu_device_id
expression = FunctionExpression(mock_function, name="test")
input_batch = Batch(frames=pd.DataFrame())
expression.evaluate(input_batch)
mock_function.to_device.assert_called_with(gpu_device_id)
gpu_mock_function.assert_called()