Python oyaml.safe_load() Examples

def test__cli__command_parse_serialize_from_stdin(serialize):
    """Check that the parser serialized output option is working.

    Not going to test for the content of the output as that is subject to change.
    """
    result = invoke_assert_code(
        args=[parse, ('-', '--format', serialize)],
        cli_input='select * from tbl',
    )
    if serialize == 'json':
        result = json.loads(result.output)
    elif serialize == 'yaml':
        result = yaml.safe_load(result.output)
    else:
        raise Exception
    result = result[0]  # only one file
    assert result['filepath'] == 'stdin' 

def test__cli__command_lint_serialize_multiple_files(serialize):
    """Check the general format of JSON output for multiple files."""
    fpath = 'test/fixtures/linter/indentation_errors.sql'

    # note the file is in here twice. two files = two payloads.
    result = invoke_assert_code(
        args=[lint, (fpath, fpath, '--format', serialize)],
        ret_code=65,
    )

    if serialize == 'json':
        result = json.loads(result.output)
    elif serialize == 'yaml':
        result = yaml.safe_load(result.output)
    else:
        raise Exception
    assert len(result) == 2 

def raw_conf(mapchete_file):
    """
    Loads a mapchete_file into a dictionary.

    Parameters
    ----------
    mapchete_file : str
        Path to a Mapchete file.

    Returns
    -------
    dictionary
    """
    if isinstance(mapchete_file, dict):
        return _map_to_new_config(mapchete_file)
    else:
        return _map_to_new_config(yaml.safe_load(open(mapchete_file, "r").read())) 

def _config_to_dict(input_config):
    if isinstance(input_config, dict):
        if "config_dir" not in input_config:
            raise MapcheteConfigError("config_dir parameter missing")
        return OrderedDict(input_config, mapchete_file=None)
    # from Mapchete file
    elif os.path.splitext(input_config)[1] == ".mapchete":
        with open(input_config, "r") as config_file:
            return OrderedDict(
                yaml.safe_load(config_file.read()),
                config_dir=os.path.dirname(os.path.realpath(input_config)),
                mapchete_file=input_config
            )
    # throw error if unknown object
    else:
        raise MapcheteConfigError(
            "Configuration has to be a dictionary or a .mapchete file.") 

def read_yaml(self, filename, render_variables=False):
        return yaml.safe_load(self.read_file(filename, render_variables)) 

def read_yaml(filename):
    with open(filename) as yaml_file:
        data = yaml_file.read()
    return yaml.safe_load(data) 

def _validate_recommended_data(data: OrderedDict) -> bool:
    """
    Ensure the returned data has the following structure:
    links:
      - name: Global Protect Skillets
        link: https://github.com/PaloAltoNetworks/GPSkillets
        branch: 90dev
        description: Configuration templates for GlobalProtect mobile users
    :param data: loaded object from oyaml.safe_load
    :return: boolean
    """

    if 'links' not in data:
        print('No links key in data')
        return False

    if type(data['links']) is not list:
        print('links is not a valid list!')
        return False

    for link in data['links']:
        if type(link) is not OrderedDict and type(link) is not dict:
            print('link entry is not a dict')
            return False

        if 'name' not in link or 'link' not in link or 'description' not in link:
            print('link entry does not have all required keys')
            return False

    return True 

def get_config():
    config_file = get_user_config_path()

    with open(config_file) as f:
        return yaml.safe_load(f) 

def get_default_config():
    data_dir = os.path.dirname(os.path.realpath(__file__))
    data_path = os.path.join(data_dir, "data/shell-config.yaml")

    with open(data_path) as f:
        return yaml.safe_load(f) 

def test_validate_config_valid(mock_exit, mock_print, test_input, expected):
    with open("tests/data/test-config.yaml") as f:
        config = yaml.safe_load(f)
        returned_config = kafkashell.config.validate_config(config)

        assert not mock_print.called
        assert not mock_exit.called
        assert config == returned_config 

def test_validate_config_is_valid_with_environment_variables(mock_exit, mock_print, test_input, expected):
    with open("tests/data/test-environment-variables-config.yaml") as f:
        config = yaml.safe_load(f)
        returned_config = kafkashell.config.validate_config(config)

        assert not mock_print.called
        assert not mock_exit.called
        assert config == returned_config 

def setup_settings_with_real_completer_for_test(mock_settings):
    with open("kafkashell/data/completer.json") as f:
        with open("kafkashell/data/shell-config.yaml") as fc:
            settings = mock_settings.return_value
            settings.selected_cluster = "local"
            settings.enable_help = True
            settings.enable_fuzzy_search = True
            settings.user_config = yaml.safe_load(fc)
            settings.commands = json.load(f)["commands"]
            return settings 

def get_test_config(config_name="test"):
    data_dir = os.path.dirname(os.path.realpath(__file__))
    data_path = os.path.realpath(os.path.join(data_dir, "../tests/data/{0}-config.yaml".format(config_name)))
    with open(data_path) as f:
        return yaml.safe_load(f) 

def load_yaml(fpath):
    """Load a yaml structure and process it into a tuple."""
    # Load raw file
    with open(fpath) as f:
        raw = f.read()
    # Parse the yaml
    obj = oyaml.safe_load(raw)
    # Return the parsed and structured object
    return process_struct(obj)[0] 

def test__cli__command_lint_serialize_from_stdin(serialize, sql, expected, exit_code):
    """Check an explicit serialized return value for a single error."""
    result = invoke_assert_code(
        args=[lint, ('-', '--rules', 'L010', '--format', serialize)],
        cli_input=sql,
        ret_code=exit_code
    )

    if serialize == 'json':
        assert json.loads(result.output) == expected
    elif serialize == 'yaml':
        assert yaml.safe_load(result.output) == expected
    else:
        raise Exception 

def test_settings(mock_init_config, mock_init_history, mock_get_config, mock_get_completer, mock_save_config):
    with open("tests/data/test-config.yaml") as f:
        with open("tests/data/test-modified-config.yaml") as fm:
            config_json = yaml.safe_load(f)
            modified_json = yaml.safe_load(fm)
            mock_get_completer.return_value = {"commands": {}}
            mock_get_config.return_value = config_json

            # test init
            settings = kafkashell.settings.Settings()

            mock_init_config.assert_called_once()
            mock_init_history.assert_called_once()
            mock_get_config.assert_called_once()
            mock_get_completer.assert_called_once()
            assert settings.enable_help is True
            assert settings.enable_auto_complete is True
            assert settings.enable_auto_suggest is True
            assert settings.cluster == "local"

            # test set_enable_help
            settings.set_enable_help(False)
            assert settings.enable_help is False

            # test set_enable_fuzzy_search
            settings.set_enable_fuzzy_search(False)
            assert settings.enable_fuzzy_search is False

            # test set_next_cluster & get_cluster_details
            settings.set_next_cluster()
            assert settings.cluster == "test"
            assert settings.get_cluster_details() == config_json["clusters"]["test"]

            # test save_settings
            settings.save_settings()
            mock_save_config.assert_called_once_with(modified_json)

            # test save_settings when enableSaveOnExit is false
            mock_save_config.reset_mock()
            settings.enable_save_on_exit = False
            settings.save_settings()
            assert not mock_save_config.called

            # test things can change back
            settings.set_enable_help(True)
            assert settings.enable_help is True

            settings.set_enable_fuzzy_search(True)
            assert settings.enable_fuzzy_search is True

            settings.set_next_cluster()
            assert settings.cluster == "local"
            assert settings.get_cluster_details() == config_json["clusters"]["local"]