borgmatic/tests/unit/actions/test_restore.py

import collections

import pytest
from flexmock import flexmock

import borgmatic.actions.restore as module


def test_get_configured_data_source_matches_data_source_by_name():
    flexmock(module).should_receive('dumps_match').and_return(False)
    flexmock(module).should_receive('dumps_match').with_args(
        module.Dump('postgresql_databases', 'bar'),
        module.Dump('postgresql_databases', 'bar'),
    ).and_return(True)

    assert module.get_configured_data_source(
        config={
            'other_databases': [{'name': 'other'}],
            'postgresql_databases': [{'name': 'foo'}, {'name': 'bar'}],
        },
        restore_dump=module.Dump('postgresql_databases', 'bar'),
    ) == ('postgresql_databases', {'name': 'bar'})


def test_get_configured_data_source_matches_nothing_when_nothing_configured():
    flexmock(module).should_receive('dumps_match').and_return(False)

    assert module.get_configured_data_source(
        config={},
        restore_dump=module.Dump('postgresql_databases', 'quux'),
    ) == (None, None)


def test_get_configured_data_source_matches_nothing_when_data_source_name_not_configured():
    flexmock(module).should_receive('dumps_match').and_return(False)

    assert module.get_configured_data_source(
        config={
            'postgresql_databases': [{'name': 'foo'}],
        },
        restore_dump=module.Dump('postgresql_databases', 'quux'),
    ) == (None, None)


def test_get_configured_data_source_with_unspecified_hook_matches_data_source_by_name():
    flexmock(module).should_receive('dumps_match').and_return(False)
    flexmock(module).should_receive('dumps_match').with_args(
        module.Dump('postgresql_databases', 'bar'),
        module.Dump(module.UNSPECIFIED, 'bar'),
    ).and_return(True)

    assert module.get_configured_data_source(
        config={
            'other_databases': [{'name': 'other'}],
            'postgresql_databases': [{'name': 'foo'}, {'name': 'bar'}],
        },
        restore_dump=module.Dump(module.UNSPECIFIED, 'bar'),
    ) == ('postgresql_databases', {'name': 'bar'})


def test_strip_path_prefix_from_extracted_dump_destination_renames_first_matching_databases_subdirectory():
    flexmock(module.os).should_receive('walk').and_return(
        [
            ('/foo', flexmock(), flexmock()),
            ('/foo/bar', flexmock(), flexmock()),
            ('/foo/bar/postgresql_databases', flexmock(), flexmock()),
            ('/foo/bar/mariadb_databases', flexmock(), flexmock()),
        ]
    )

    flexmock(module.shutil).should_receive('move').with_args(
        '/foo/bar/postgresql_databases', '/run/user/0/borgmatic/postgresql_databases'
    ).once()
    flexmock(module.shutil).should_receive('move').with_args(
        '/foo/bar/mariadb_databases', '/run/user/0/borgmatic/mariadb_databases'
    ).never()

    module.strip_path_prefix_from_extracted_dump_destination('/foo', '/run/user/0/borgmatic')


def test_restore_single_dump_extracts_and_restores_single_file_dump():
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hooks').with_args(
        'make_data_source_dump_patterns', object, object, object, object, object
    ).and_return({'postgresql': flexmock()})
    flexmock(module.tempfile).should_receive('mkdtemp').never()
    flexmock(module.borgmatic.hooks.data_source.dump).should_receive(
        'convert_glob_patterns_to_borg_pattern'
    ).and_return(flexmock())
    flexmock(module.borgmatic.borg.extract).should_receive('extract_archive').and_return(
        flexmock()
    ).once()
    flexmock(module).should_receive('strip_path_prefix_from_extracted_dump_destination').never()
    flexmock(module.shutil).should_receive('rmtree').never()
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hook').with_args(
        function_name='restore_data_source_dump',
        config=object,
        log_prefix=object,
        hook_name=object,
        data_source=object,
        dry_run=object,
        extract_process=object,
        connection_params=object,
        borgmatic_runtime_directory=object,
    ).once()

    module.restore_single_dump(
        repository={'path': 'test.borg'},
        config=flexmock(),
        local_borg_version=flexmock(),
        global_arguments=flexmock(dry_run=False),
        local_path=None,
        remote_path=None,
        archive_name=flexmock(),
        hook_name='postgresql',
        data_source={'name': 'test', 'format': 'plain'},
        connection_params=flexmock(),
        borgmatic_runtime_directory='/run/borgmatic',
    )


def test_restore_single_dump_extracts_and_restores_directory_dump():
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hooks').with_args(
        'make_data_source_dump_patterns', object, object, object, object, object
    ).and_return({'postgresql': flexmock()})
    flexmock(module.tempfile).should_receive('mkdtemp').once().and_return(
        '/run/user/0/borgmatic/tmp1234'
    )
    flexmock(module.borgmatic.hooks.data_source.dump).should_receive(
        'convert_glob_patterns_to_borg_pattern'
    ).and_return(flexmock())
    flexmock(module.borgmatic.borg.extract).should_receive('extract_archive').and_return(
        flexmock()
    ).once()
    flexmock(module).should_receive('strip_path_prefix_from_extracted_dump_destination').once()
    flexmock(module.shutil).should_receive('rmtree').once()
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hook').with_args(
        function_name='restore_data_source_dump',
        config=object,
        log_prefix=object,
        hook_name=object,
        data_source=object,
        dry_run=object,
        extract_process=object,
        connection_params=object,
        borgmatic_runtime_directory='/run/borgmatic',
    ).once()

    module.restore_single_dump(
        repository={'path': 'test.borg'},
        config=flexmock(),
        local_borg_version=flexmock(),
        global_arguments=flexmock(dry_run=False),
        local_path=None,
        remote_path=None,
        archive_name=flexmock(),
        hook_name='postgresql',
        data_source={'name': 'test', 'format': 'directory'},
        connection_params=flexmock(),
        borgmatic_runtime_directory='/run/borgmatic',
    )


def test_restore_single_dump_with_directory_dump_error_cleans_up_temporary_directory():
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hooks').with_args(
        'make_data_source_dump_patterns', object, object, object, object, object
    ).and_return({'postgresql': flexmock()})
    flexmock(module.tempfile).should_receive('mkdtemp').once().and_return(
        '/run/user/0/borgmatic/tmp1234'
    )
    flexmock(module.borgmatic.hooks.data_source.dump).should_receive(
        'convert_glob_patterns_to_borg_pattern'
    ).and_return(flexmock())
    flexmock(module.borgmatic.borg.extract).should_receive('extract_archive').and_raise(
        ValueError
    ).once()
    flexmock(module).should_receive('strip_path_prefix_from_extracted_dump_destination').never()
    flexmock(module.shutil).should_receive('rmtree').once()
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hook').with_args(
        function_name='restore_data_source_dump',
        config=object,
        log_prefix=object,
        hook_name=object,
        data_source=object,
        dry_run=object,
        extract_process=object,
        connection_params=object,
        borgmatic_runtime_directory='/run/user/0/borgmatic/tmp1234',
    ).never()

    with pytest.raises(ValueError):
        module.restore_single_dump(
            repository={'path': 'test.borg'},
            config=flexmock(),
            local_borg_version=flexmock(),
            global_arguments=flexmock(dry_run=False),
            local_path=None,
            remote_path=None,
            archive_name=flexmock(),
            hook_name='postgresql',
            data_source={'name': 'test', 'format': 'directory'},
            connection_params=flexmock(),
            borgmatic_runtime_directory='/run/borgmatic',
        )


def test_restore_single_dump_with_directory_dump_and_dry_run_skips_directory_move_and_cleanup():
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hooks').with_args(
        'make_data_source_dump_patterns', object, object, object, object, object
    ).and_return({'postgresql': flexmock()})
    flexmock(module.tempfile).should_receive('mkdtemp').once().and_return('/run/borgmatic/tmp1234')
    flexmock(module.borgmatic.hooks.data_source.dump).should_receive(
        'convert_glob_patterns_to_borg_pattern'
    ).and_return(flexmock())
    flexmock(module.borgmatic.borg.extract).should_receive('extract_archive').and_return(
        flexmock()
    ).once()
    flexmock(module).should_receive('strip_path_prefix_from_extracted_dump_destination').never()
    flexmock(module.shutil).should_receive('rmtree').never()
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hook').with_args(
        function_name='restore_data_source_dump',
        config=object,
        log_prefix=object,
        hook_name=object,
        data_source=object,
        dry_run=object,
        extract_process=object,
        connection_params=object,
        borgmatic_runtime_directory='/run/borgmatic',
    ).once()

    module.restore_single_dump(
        repository={'path': 'test.borg'},
        config=flexmock(),
        local_borg_version=flexmock(),
        global_arguments=flexmock(dry_run=True),
        local_path=None,
        remote_path=None,
        archive_name=flexmock(),
        hook_name='postgresql',
        data_source={'name': 'test', 'format': 'directory'},
        connection_params=flexmock(),
        borgmatic_runtime_directory='/run/borgmatic',
    )


def test_collect_dumps_from_archive_parses_archive_paths():
    flexmock(module.borgmatic.config.paths).should_receive(
        'get_borgmatic_source_directory'
    ).and_return('/root/.borgmatic')
    flexmock(module.borgmatic.hooks.data_source.dump).should_receive(
        'make_data_source_dump_path'
    ).and_return('')
    flexmock(module.borgmatic.borg.list).should_receive('capture_archive_listing').and_return(
        [
            'borgmatic/postgresql_databases/localhost/foo',
            'borgmatic/postgresql_databases/localhost/bar',
            'borgmatic/mysql_databases/localhost/quux',
        ]
    )

    archive_dumps = module.collect_dumps_from_archive(
        repository={'path': 'repo'},
        archive='archive',
        config={},
        local_borg_version=flexmock(),
        global_arguments=flexmock(log_json=False),
        local_path=flexmock(),
        remote_path=flexmock(),
        borgmatic_runtime_directory='/run/borgmatic',
    )

    assert archive_dumps == {
        module.Dump('postgresql_databases', 'foo'),
        module.Dump('postgresql_databases', 'bar'),
        module.Dump('mysql_databases', 'quux'),
    }


def test_collect_dumps_from_archive_parses_archive_paths_with_different_base_directories():
    flexmock(module.borgmatic.config.paths).should_receive(
        'get_borgmatic_source_directory'
    ).and_return('/root/.borgmatic')
    flexmock(module.borgmatic.hooks.data_source.dump).should_receive(
        'make_data_source_dump_path'
    ).and_return('')
    flexmock(module.borgmatic.borg.list).should_receive('capture_archive_listing').and_return(
        [
            'borgmatic/postgresql_databases/localhost/foo',
            '.borgmatic/postgresql_databases/localhost/bar',
            '/root/.borgmatic/postgresql_databases/localhost/baz',
            '/var/run/0/borgmatic/mysql_databases/localhost/quux',
        ]
    )

    archive_dumps = module.collect_dumps_from_archive(
        repository={'path': 'repo'},
        archive='archive',
        config={},
        local_borg_version=flexmock(),
        global_arguments=flexmock(log_json=False),
        local_path=flexmock(),
        remote_path=flexmock(),
        borgmatic_runtime_directory='/run/borgmatic',
    )

    assert archive_dumps == {
        module.Dump('postgresql_databases', 'foo'),
        module.Dump('postgresql_databases', 'bar'),
        module.Dump('postgresql_databases', 'baz'),
        module.Dump('mysql_databases', 'quux'),
    }


def test_collect_dumps_from_archive_parses_directory_format_archive_paths():
    flexmock(module.borgmatic.config.paths).should_receive(
        'get_borgmatic_source_directory'
    ).and_return('/root/.borgmatic')
    flexmock(module.borgmatic.hooks.data_source.dump).should_receive(
        'make_data_source_dump_path'
    ).and_return('')
    flexmock(module.borgmatic.borg.list).should_receive('capture_archive_listing').and_return(
        [
            'borgmatic/postgresql_databases/localhost/foo/table1',
            'borgmatic/postgresql_databases/localhost/foo/table2',
        ]
    )

    archive_dumps = module.collect_dumps_from_archive(
        repository={'path': 'repo'},
        archive='archive',
        config={},
        local_borg_version=flexmock(),
        global_arguments=flexmock(log_json=False),
        local_path=flexmock(),
        remote_path=flexmock(),
        borgmatic_runtime_directory='/run/borgmatic',
    )

    assert archive_dumps == {
        module.Dump('postgresql_databases', 'foo'),
    }


def test_collect_dumps_from_archive_skips_bad_archive_paths():
    flexmock(module.borgmatic.config.paths).should_receive(
        'get_borgmatic_source_directory'
    ).and_return('/root/.borgmatic')
    flexmock(module.borgmatic.hooks.data_source.dump).should_receive(
        'make_data_source_dump_path'
    ).and_return('')
    flexmock(module.borgmatic.borg.list).should_receive('capture_archive_listing').and_return(
        [
            'borgmatic/postgresql_databases/localhost/foo',
            'borgmatic/invalid',
            'invalid/as/well',
            '',
        ]
    )

    archive_dumps = module.collect_dumps_from_archive(
        repository={'path': 'repo'},
        archive='archive',
        config={},
        local_borg_version=flexmock(),
        global_arguments=flexmock(log_json=False),
        local_path=flexmock(),
        remote_path=flexmock(),
        borgmatic_runtime_directory='/run/borgmatic',
    )

    assert archive_dumps == {
        module.Dump('postgresql_databases', 'foo'),
    }


def test_get_dumps_to_restore_gets_requested_dumps_found_in_archive():
    dumps_from_archive = {
        module.Dump('postgresql_databases', 'foo'),
        module.Dump('postgresql_databases', 'bar'),
        module.Dump('postgresql_databases', 'baz'),
    }
    flexmock(module).should_receive('dumps_match').and_return(False)
    flexmock(module).should_receive('dumps_match').with_args(
        module.Dump(module.UNSPECIFIED, 'foo'),
        module.Dump('postgresql_databases', 'foo'),
    ).and_return(True)
    flexmock(module).should_receive('dumps_match').with_args(
        module.Dump(module.UNSPECIFIED, 'bar'),
        module.Dump('postgresql_databases', 'bar'),
    ).and_return(True)

    assert module.get_dumps_to_restore(
        restore_arguments=flexmock(
            hook=None,
            data_sources=['foo', 'bar'],
            original_hostname=None,
            original_port=None,
        ),
        dumps_from_archive=dumps_from_archive,
    ) == {
        module.Dump('postgresql_databases', 'foo'),
        module.Dump('postgresql_databases', 'bar'),
    }


def test_get_dumps_to_restore_raises_for_requested_dumps_missing_from_archive():
    dumps_from_archive = {
        module.Dump('postgresql_databases', 'foo'),
    }
    flexmock(module).should_receive('dumps_match').and_return(False)

    with pytest.raises(ValueError):
        module.get_dumps_to_restore(
            restore_arguments=flexmock(
                hook=None,
                data_sources=['foo', 'bar'],
                original_hostname=None,
                original_port=None,
            ),
            dumps_from_archive=dumps_from_archive,
        )


def test_get_dumps_to_restore_without_requested_dumps_finds_all_archive_dumps():
    dumps_from_archive = {
        module.Dump('postgresql_databases', 'foo'),
        module.Dump('postgresql_databases', 'bar'),
    }
    flexmock(module).should_receive('dumps_match').and_return(False)

    assert module.get_dumps_to_restore(
        restore_arguments=flexmock(
            hook=None,
            data_sources=[],
            original_hostname=None,
            original_port=None,
        ),
        dumps_from_archive=dumps_from_archive,
    ) == dumps_from_archive


def test_get_dumps_to_restore_with_all_in_requested_dumps_finds_all_archive_dumps():
    dumps_from_archive = {
        module.Dump('postgresql_databases', 'foo'),
        module.Dump('postgresql_databases', 'bar'),
    }
    flexmock(module).should_receive('dumps_match').and_return(False)
    flexmock(module).should_receive('dumps_match').with_args(
        module.Dump(module.UNSPECIFIED, 'foo'),
        module.Dump('postgresql_databases', 'foo'),
    ).and_return(True)
    flexmock(module).should_receive('dumps_match').with_args(
        module.Dump(module.UNSPECIFIED, 'bar'),
        module.Dump('postgresql_databases', 'bar'),
    ).and_return(True)

    assert module.get_dumps_to_restore(
        restore_arguments=flexmock(
            hook=None,
            data_sources=['all'],
            original_hostname=None,
            original_port=None,
        ),
        dumps_from_archive=dumps_from_archive,
    ) == dumps_from_archive


def test_get_dumps_to_restore_with_all_in_requested_dumps_plus_additional_requested_dumps_omits_duplicates():
    dumps_from_archive = {
        module.Dump('postgresql_databases', 'foo'),
        module.Dump('postgresql_databases', 'bar'),
    }
    flexmock(module).should_receive('dumps_match').and_return(False)
    flexmock(module).should_receive('dumps_match').with_args(
        module.Dump(module.UNSPECIFIED, 'foo'),
        module.Dump('postgresql_databases', 'foo'),
    ).and_return(True)
    flexmock(module).should_receive('dumps_match').with_args(
        module.Dump(module.UNSPECIFIED, 'bar'),
        module.Dump('postgresql_databases', 'bar'),
    ).and_return(True)

    assert (
        module.get_dumps_to_restore(
            restore_arguments=flexmock(
                hook=None,
                data_sources=['all', 'foo', 'bar'],
                original_hostname=None,
                original_port=None,
            ),
            dumps_from_archive=dumps_from_archive,
        )
        == dumps_from_archive
    )


def test_get_dumps_to_restore_raises_for_all_in_requested_dumps_and_requested_dumps_missing_from_archives():
    flexmock(module).should_receive('dumps_match').and_return(False)
    flexmock(module).should_receive('dumps_match').with_args(
        module.Dump(module.UNSPECIFIED, 'foo'),
        module.Dump('postgresql_databases', 'foo'),
    ).and_return(True)

    with pytest.raises(ValueError):
        module.get_dumps_to_restore(
            restore_arguments=flexmock(
                hook=None,
                data_sources=['all', 'foo', 'bar'],
                original_hostname=None,
                original_port=None,
            ),
            dumps_from_archive={module.Dump('postresql_databases', 'foo')},
        )


def test_ensure_requested_dumps_restored_with_all_dumps_restored_does_not_raise():
    module.ensure_requested_dumps_restored(
        dumps_to_restore={
            module.Dump(hook_name='postgresql_databases', data_source_name='foo'),
            module.Dump(hook_name='postgresql_databases', data_source_name='bar'),
        },
        dumps_actually_restored={
            module.Dump(hook_name='postgresql_databases', data_source_name='foo'),
            module.Dump(hook_name='postgresql_databases', data_source_name='bar'),
        },
    )


def test_ensure_requested_dumps_restored_with_no_dumps_raises():
    with pytest.raises(ValueError):
        module.ensure_requested_dumps_restored(
            dumps_to_restore={},
            dumps_actually_restored={},
        )


def test_ensure_requested_dumps_restored_with_missing_dumps_raises():
    with pytest.raises(ValueError):
        module.ensure_requested_dumps_restored(
            dumps_to_restore={module.Dump(hook_name='postgresql_databases', data_source_name='foo')},
            dumps_actually_restored={
                module.Dump(hook_name='postgresql_databases', data_source_name='bar')
            },
        )


def test_run_restore_restores_each_data_source():
    dumps_to_restore = {
        module.Dump(hook_name='postgresql_databases', data_source_name='foo'),
        module.Dump(hook_name='postgresql_databases', data_source_name='bar'),
    }

    flexmock(module.borgmatic.config.validate).should_receive('repositories_match').and_return(True)
    borgmatic_runtime_directory = flexmock()
    flexmock(module.borgmatic.config.paths).should_receive('Runtime_directory').and_return(
        borgmatic_runtime_directory
    )
    flexmock(module.borgmatic.config.paths).should_receive(
        'make_runtime_directory_glob'
    ).replace_with(lambda path: path)
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hooks_even_if_unconfigured')
    flexmock(module.borgmatic.borg.repo_list).should_receive('resolve_archive_name').and_return(
        flexmock()
    )
    flexmock(module).should_receive('collect_dumps_from_archive').and_return(flexmock())
    flexmock(module).should_receive('get_dumps_to_restore').and_return(dumps_to_restore)
    flexmock(module).should_receive('get_configured_data_source').and_return(
        ('postgresql_databases', {'name': 'foo'})
    ).and_return(('postgresql_databases', {'name': 'bar'}))
    flexmock(module).should_receive('restore_single_dump').with_args(
        repository=object,
        config=object,
        local_borg_version=object,
        global_arguments=object,
        local_path=object,
        remote_path=object,
        archive_name=object,
        hook_name='postgresql_databases',
        data_source={'name': 'foo', 'schemas': None},
        connection_params=object,
        borgmatic_runtime_directory=borgmatic_runtime_directory,
    ).once()
    flexmock(module).should_receive('restore_single_dump').with_args(
        repository=object,
        config=object,
        local_borg_version=object,
        global_arguments=object,
        local_path=object,
        remote_path=object,
        archive_name=object,
        hook_name='postgresql_databases',
        data_source={'name': 'bar', 'schemas': None},
        connection_params=object,
        borgmatic_runtime_directory=borgmatic_runtime_directory,
    ).once()
    flexmock(module).should_receive('ensure_requested_dumps_restored')

    module.run_restore(
        repository={'path': 'repo'},
        config=flexmock(),
        local_borg_version=flexmock(),
        restore_arguments=flexmock(
            repository='repo',
            archive='archive',
            data_sources=flexmock(),
            schemas=None,
            hostname=None,
            port=None,
            username=None,
            password=None,
            restore_path=None,
        ),
        global_arguments=flexmock(dry_run=False),
        local_path=flexmock(),
        remote_path=flexmock(),
    )


def test_run_restore_bails_for_non_matching_repository():
    flexmock(module.borgmatic.config.validate).should_receive('repositories_match').and_return(
        False
    )
    flexmock(module.borgmatic.config.paths).should_receive('Runtime_directory').and_return(
        flexmock()
    )
    flexmock(module.borgmatic.config.paths).should_receive(
        'make_runtime_directory_glob'
    ).replace_with(lambda path: path)
    flexmock(module.borgmatic.hooks.dispatch).should_receive(
        'call_hooks_even_if_unconfigured'
    ).never()
    flexmock(module).should_receive('restore_single_dump').never()

    module.run_restore(
        repository={'path': 'repo'},
        config=flexmock(),
        local_borg_version=flexmock(),
        restore_arguments=flexmock(repository='repo', archive='archive', data_sources=flexmock()),
        global_arguments=flexmock(dry_run=False),
        local_path=flexmock(),
        remote_path=flexmock(),
    )


def test_run_restore_restores_data_source_configured_with_all_name():
    dumps_to_restore = {
        module.Dump(hook_name='postgresql_databases', data_source_name='foo'),
        module.Dump(hook_name='postgresql_databases', data_source_name='bar'),
    }

    flexmock(module.borgmatic.config.validate).should_receive('repositories_match').and_return(True)
    borgmatic_runtime_directory = flexmock()
    flexmock(module.borgmatic.config.paths).should_receive('Runtime_directory').and_return(
        borgmatic_runtime_directory
    )
    flexmock(module.borgmatic.config.paths).should_receive(
        'make_runtime_directory_glob'
    ).replace_with(lambda path: path)
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hooks_even_if_unconfigured')
    flexmock(module.borgmatic.borg.repo_list).should_receive('resolve_archive_name').and_return(
        flexmock()
    )
    flexmock(module).should_receive('collect_dumps_from_archive').and_return(flexmock())
    flexmock(module).should_receive('get_dumps_to_restore').and_return(dumps_to_restore)
    flexmock(module).should_receive('get_configured_data_source').with_args(
        config=object,
        restore_dump=module.Dump(hook_name='postgresql_databases', data_source_name='foo'),
    ).and_return(('postgresql_databases', {'name': 'foo'}))
    flexmock(module).should_receive('get_configured_data_source').with_args(
        config=object,
        restore_dump=module.Dump(hook_name='postgresql_databases', data_source_name='bar'),
    ).and_return((None, None))
    flexmock(module).should_receive('get_configured_data_source').with_args(
        config=object,
        restore_dump=module.Dump(hook_name='postgresql_databases', data_source_name='all'),
    ).and_return(('postgresql_databases', {'name': 'bar'}))
    flexmock(module).should_receive('restore_single_dump').with_args(
        repository=object,
        config=object,
        local_borg_version=object,
        global_arguments=object,
        local_path=object,
        remote_path=object,
        archive_name=object,
        hook_name='postgresql_databases',
        data_source={'name': 'foo', 'schemas': None},
        connection_params=object,
        borgmatic_runtime_directory=borgmatic_runtime_directory,
    ).once()
    flexmock(module).should_receive('restore_single_dump').with_args(
        repository=object,
        config=object,
        local_borg_version=object,
        global_arguments=object,
        local_path=object,
        remote_path=object,
        archive_name=object,
        hook_name='postgresql_databases',
        data_source={'name': 'bar', 'schemas': None},
        connection_params=object,
        borgmatic_runtime_directory=borgmatic_runtime_directory,
    ).once()
    flexmock(module).should_receive('ensure_requested_dumps_restored')

    module.run_restore(
        repository={'path': 'repo'},
        config=flexmock(),
        local_borg_version=flexmock(),
        restore_arguments=flexmock(
            repository='repo',
            archive='archive',
            data_sources=flexmock(),
            schemas=None,
            hostname=None,
            port=None,
            username=None,
            password=None,
            restore_path=None,
        ),
        global_arguments=flexmock(dry_run=False),
        local_path=flexmock(),
        remote_path=flexmock(),
    )


def test_run_restore_skips_missing_data_source():
    dumps_to_restore = {
        module.Dump(hook_name='postgresql_databases', data_source_name='foo'),
        module.Dump(hook_name='postgresql_databases', data_source_name='bar'),
    }

    flexmock(module.borgmatic.config.validate).should_receive('repositories_match').and_return(True)
    borgmatic_runtime_directory = flexmock()
    flexmock(module.borgmatic.config.paths).should_receive('Runtime_directory').and_return(
        borgmatic_runtime_directory
    )
    flexmock(module.borgmatic.config.paths).should_receive(
        'make_runtime_directory_glob'
    ).replace_with(lambda path: path)
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hooks_even_if_unconfigured')
    flexmock(module.borgmatic.borg.repo_list).should_receive('resolve_archive_name').and_return(
        flexmock()
    )
    flexmock(module).should_receive('collect_dumps_from_archive').and_return(flexmock())
    flexmock(module).should_receive('get_dumps_to_restore').and_return(dumps_to_restore)
    flexmock(module).should_receive('get_configured_data_source').with_args(
        config=object,
        restore_dump=module.Dump(hook_name='postgresql_databases', data_source_name='foo'),
    ).and_return(('postgresql_databases', {'name': 'foo'}))
    flexmock(module).should_receive('get_configured_data_source').with_args(
        config=object,
        restore_dump=module.Dump(hook_name='postgresql_databases', data_source_name='bar'),
    ).and_return((None, None))
    flexmock(module).should_receive('get_configured_data_source').with_args(
        config=object,
        restore_dump=module.Dump(hook_name='postgresql_databases', data_source_name='all'),
    ).and_return((None, None))
    flexmock(module).should_receive('restore_single_dump').with_args(
        repository=object,
        config=object,
        local_borg_version=object,
        global_arguments=object,
        local_path=object,
        remote_path=object,
        archive_name=object,
        hook_name='postgresql_databases',
        data_source={'name': 'foo', 'schemas': None},
        connection_params=object,
        borgmatic_runtime_directory=borgmatic_runtime_directory,
    ).once()
    flexmock(module).should_receive('restore_single_dump').with_args(
        repository=object,
        config=object,
        local_borg_version=object,
        global_arguments=object,
        local_path=object,
        remote_path=object,
        archive_name=object,
        hook_name='postgresql_databases',
        data_source={'name': 'bar', 'schemas': None},
        connection_params=object,
        borgmatic_runtime_directory=borgmatic_runtime_directory,
    ).never()
    flexmock(module).should_receive('ensure_requested_dumps_restored')

    module.run_restore(
        repository={'path': 'repo'},
        config=flexmock(),
        local_borg_version=flexmock(),
        restore_arguments=flexmock(
            repository='repo',
            archive='archive',
            data_sources=flexmock(),
            schemas=None,
            hostname=None,
            port=None,
            username=None,
            password=None,
            restore_path=None,
        ),
        global_arguments=flexmock(dry_run=False),
        local_path=flexmock(),
        remote_path=flexmock(),
    )


def test_run_restore_restores_data_sources_from_different_hooks():
    dumps_to_restore = {
        module.Dump(hook_name='postgresql_databases', data_source_name='foo'),
        module.Dump(hook_name='mysql_databases', data_source_name='foo'),
    }

    flexmock(module.borgmatic.config.validate).should_receive('repositories_match').and_return(True)
    borgmatic_runtime_directory = flexmock()
    flexmock(module.borgmatic.config.paths).should_receive('Runtime_directory').and_return(
        borgmatic_runtime_directory
    )
    flexmock(module.borgmatic.config.paths).should_receive(
        'make_runtime_directory_glob'
    ).replace_with(lambda path: path)
    flexmock(module.borgmatic.hooks.dispatch).should_receive('call_hooks_even_if_unconfigured')
    flexmock(module.borgmatic.borg.repo_list).should_receive('resolve_archive_name').and_return(
        flexmock()
    )
    flexmock(module).should_receive('collect_dumps_from_archive').and_return(flexmock())
    flexmock(module).should_receive('get_dumps_to_restore').and_return(dumps_to_restore)
    flexmock(module).should_receive('get_configured_data_source').with_args(
        config=object,
        restore_dump=module.Dump(hook_name='postgresql_databases', data_source_name='foo'),
    ).and_return(('postgresql_databases', {'name': 'foo'}))
    flexmock(module).should_receive('get_configured_data_source').with_args(
        config=object,
        restore_dump=module.Dump(hook_name='mysql_databases', data_source_name='foo'),
    ).and_return(('mysql_databases', {'name': 'bar'}))
    flexmock(module).should_receive('restore_single_dump').with_args(
        repository=object,
        config=object,
        local_borg_version=object,
        global_arguments=object,
        local_path=object,
        remote_path=object,
        archive_name=object,
        hook_name='postgresql_databases',
        data_source={'name': 'foo', 'schemas': None},
        connection_params=object,
        borgmatic_runtime_directory=borgmatic_runtime_directory,
    ).once()
    flexmock(module).should_receive('restore_single_dump').with_args(
        repository=object,
        config=object,
        local_borg_version=object,
        global_arguments=object,
        local_path=object,
        remote_path=object,
        archive_name=object,
        hook_name='mysql_databases',
        data_source={'name': 'bar', 'schemas': None},
        connection_params=object,
        borgmatic_runtime_directory=borgmatic_runtime_directory,
    ).once()
    flexmock(module).should_receive('ensure_requested_dumps_restored')

    module.run_restore(
        repository={'path': 'repo'},
        config=flexmock(),
        local_borg_version=flexmock(),
        restore_arguments=flexmock(
            repository='repo',
            archive='archive',
            data_sources=flexmock(),
            schemas=None,
            hostname=None,
            port=None,
            username=None,
            password=None,
            restore_path=None,
        ),
        global_arguments=flexmock(dry_run=False),
        local_path=flexmock(),
        remote_path=flexmock(),
    )