openfisca-run-test replacement using Pytest

test_yamls.py

from yaml import CLoader as Loader, CDumper as Dumper

import pytest
import glob
import yaml
import numpy

from openfisca_core import scenarios
from openfisca_core.scripts import build_tax_benefit_system

def load_tests(yamls, tax_benefit_system):
  return [parse_test(test, file, tax_benefit_system) for file in yamls for test in load_file(file)]

def load_file(yaml_path):
  with open(yaml_path) as yaml_file:
    contents = yaml.load(yaml_file, Loader=Loader)
    if isinstance(contents, dict):
      return [contents]
    else:
      return contents

def parse_test(test, path, tax_benefit_system):
  parsed, error = scenarios.make_json_or_python_to_test(
      tax_benefit_system = tax_benefit_system
      )(test)
  return parsed

yamls = glob.glob("tests/**/*.yaml")
@pytest.fixture(params=load_tests(yamls, build_tax_benefit_system(None, None, None)))
def yaml(request):
  return request.param

def actual(yaml):
  scenario = yaml['scenario']
  scenario.suggest()
  simulation = scenario.new_simulation(trace = False)
  output_variables = yaml.get('output_variables')
  if output_variables is not None:
    actual_results = {}
    for variable_name, expected_value in output_variables.items():
      if isinstance(expected_value, dict):
        for requested_period, expected_value_at_period in expected_value.items():
          actual_results[variable_name] = simulation.calculate(variable_name, requested_period)
        return actual_results

def expected(yaml):
  expected = {}
  output_variables = yaml.get('output_variables')
  if output_variables is not None:
    for variable_name, expected_value in output_variables.items():
      if isinstance(expected_value, dict):
        for requested_period, expected_value_at_period in expected_value.items():
          # Single values got cast into Numpy arrays earlier, undo that…
          dtype = expected_value_at_period.dtype
          if len(expected_value_at_period) == 1 and numpy.issubdtype(dtype, numpy.number):
            expected_value_at_period = expected_value_at_period[0]
          # …then proceed
          expected[variable_name] = pytest.approx(expected_value_at_period,
              rel = yaml.get('relative_error_margin'),
              abs = yaml.get('absolute_error_margin')
            )
        return expected
  return None

def test_yaml(yaml):
  assert expected(yaml) == actual(yaml)