from typing import Dict, Any

from flows.data_transformations.abstract import DataTransformation


class CorrectnessFlag(DataTransformation):
    def __init__(self, output_key, input_key):
        super().__init__(output_key)
        self.input_key = input_key

    def __call__(self, data_dict: Dict[str, Any], **kwargs) -> Dict[str, Any]:
        all_tests_passed = all([test_result["status"] for test_result in data_dict[self.input_key]])
        data_dict[self.output_key] = all_tests_passed
        return data_dict