lakehouse_engine.utils.expectations_utils

  1"""Utilities to be used by custom expectations."""
  2
  3from typing import Any, Dict, Optional
  4
  5from great_expectations.core import ExpectationConfiguration
  6from great_expectations.execution_engine import ExecutionEngine
  7from great_expectations.expectations.expectation import Expectation
  8
  9
 10def validate_result(
 11    expectation: Expectation,
 12    configuration: ExpectationConfiguration,
 13    metrics: dict,
 14    runtime_configuration: Optional[dict],
 15    execution_engine: Optional[ExecutionEngine],
 16    base_expectation: Expectation,
 17) -> Any:
 18    """Validates that the unexpected_index_list in the tests is corretly defined.
 19
 20    Additionally, it validates the expectation using the GE _validate method.
 21
 22    Args:
 23        expectation: Expectation to validate.
 24        configuration: Configuration used in the test.
 25        metrics: Test result metrics.
 26        runtime_configuration: Configuration used when running the expectation.
 27        execution_engine: Execution engine used in the expectation.
 28        base_expectation: Base expectation to validate.
 29    """
 30    example_unexpected_index_list = _get_example_unexpected_index_list(
 31        expectation, configuration
 32    )
 33
 34    test_unexpected_index_list = _get_test_unexpected_index_list(
 35        expectation.map_metric, metrics
 36    )
 37    if example_unexpected_index_list:
 38        if example_unexpected_index_list != test_unexpected_index_list:
 39            raise AssertionError(
 40                f"Example unexpected_index_list: {example_unexpected_index_list}\n"
 41                f"Test unexpected_index_list: {test_unexpected_index_list}"
 42            )
 43    return base_expectation._validate(
 44        expectation, configuration, metrics, runtime_configuration, execution_engine
 45    )
 46
 47
 48def _get_example_unexpected_index_list(
 49    expectation: Expectation, configuration: ExpectationConfiguration
 50) -> list:
 51    """Retrieves the unexpected index list defined from the example used on the test.
 52
 53    This needs to be done manually because GE allows us to get either the complete
 54    output of the test or the complete configuration used on the test.
 55    To get around this limitation this function is used to fetch the example used
 56    in the test directly from the expectation itself.
 57
 58    Args:
 59        expectation: Expectation to fetch the examples.
 60        configuration: Configuration used in the test.
 61
 62    Returns:
 63         List of unexpected indexes defined in the example used.
 64    """
 65    filtered_example: dict = {"out": {"unexpected_index_list": []}}
 66
 67    for example in expectation.examples:
 68        for test in example["tests"]:  # type: ignore
 69            example_result_format = []
 70            if "result_format" in configuration["kwargs"]:
 71                example_result_format = configuration["kwargs"]["result_format"]
 72
 73            if test["in"]["result_format"] == example_result_format:
 74                filtered_example = test
 75
 76    example_unexpected_index_list = []
 77    if "unexpected_index_list" in filtered_example["out"]:
 78        example_unexpected_index_list = filtered_example["out"]["unexpected_index_list"]
 79
 80    return example_unexpected_index_list
 81
 82
 83def _get_test_unexpected_index_list(metric_name: str, metrics: Dict) -> list:
 84    """Retrieves the unexpected index list from the test case that has been run.
 85
 86    Args:
 87        metric_name: Name of the metric to retrieve the unexpected index list.
 88        metrics: Metric values resulting from the test.
 89
 90    Returns:
 91         List of unexpected indexes retrieved form the test.
 92    """
 93    test_unexpected_index_list = []
 94    if f"{metric_name}.unexpected_index_list" in metrics:
 95        if metrics[f"{metric_name}.unexpected_index_list"]:
 96            test_unexpected_index_list = metrics[f"{metric_name}.unexpected_index_list"]
 97        else:
 98            test_unexpected_index_list = []
 99
100    return test_unexpected_index_list