# Copyright 2019 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Allow tests to assert recipe success or failure.

TODO(crbug/1010715) Remove this and use functionality in core recipe engine.
"""

from recipe_engine import recipe_test_api
from recipe_engine import post_process


class StatusCheckTestApi(recipe_test_api.RecipeTestApi):

  def test(self, name, status='success'):
    return super(StatusCheckTestApi, self).test(name) + self.status(status)

  def status(self, status='success'):
    """Returns step data to check status of recipe at end of test.

    Args:
      status: One of 'success' (default), 'failure', or 'infra_failure'. The
          result of the test case will be required to match this.
    """

    assertion_type = {
        'infra_failure': post_process.StatusException,
        'failure': post_process.StatusFailure,
        'success': post_process.StatusSuccess,
    }[status]
    return self.post_process(assertion_type)