add tests

This commit is contained in:
Divyansh Singh 2023-04-15 02:57:51 +05:30
parent 2fea429d78
commit f273e82d74
5 changed files with 122 additions and 31 deletions

View file

@ -634,7 +634,7 @@ def make_parsers():
metavar='NAME', metavar='NAME',
nargs='+', nargs='+',
dest='schemas', dest='schemas',
help="Names of schemas to restore from the database, defaults to all schemas. Schemas are only supported for PostgreSQL and MongoDB databases", help='Names of schemas to restore from the database, defaults to all schemas. Schemas are only supported for PostgreSQL and MongoDB databases',
) )
restore_group.add_argument( restore_group.add_argument(
'-h', '--help', action='help', help='Show this help message and exit' '-h', '--help', action='help', help='Show this help message and exit'

View file

@ -1,4 +1,5 @@
import csv import csv
import itertools
import logging import logging
import os import os
@ -225,11 +226,12 @@ def restore_database_dump(database_config, log_prefix, location_config, dry_run,
+ (('--username', database['username']) if 'username' in database else ()) + (('--username', database['username']) if 'username' in database else ())
+ (tuple(database['restore_options'].split(' ')) if 'restore_options' in database else ()) + (tuple(database['restore_options'].split(' ')) if 'restore_options' in database else ())
+ (() if extract_process else (dump_filename,)) + (() if extract_process else (dump_filename,))
+ tuple(
itertools.chain.from_iterable(('--schema', schema) for schema in database['schemas'])
if database['schemas']
else ()
)
) )
if database['schemas']:
for schema in database['schemas']:
restore_command += ('--schema', schema)
extra_environment = make_extra_environment(database) extra_environment = make_extra_environment(database)

View file

@ -233,7 +233,7 @@ def test_run_restore_restores_each_database():
remote_path=object, remote_path=object,
archive_name=object, archive_name=object,
hook_name='postgresql_databases', hook_name='postgresql_databases',
database={'name': 'foo'}, database={'name': 'foo', 'schemas': None},
).once() ).once()
flexmock(module).should_receive('restore_single_database').with_args( flexmock(module).should_receive('restore_single_database').with_args(
repository=object, repository=object,
@ -246,7 +246,7 @@ def test_run_restore_restores_each_database():
remote_path=object, remote_path=object,
archive_name=object, archive_name=object,
hook_name='postgresql_databases', hook_name='postgresql_databases',
database={'name': 'bar'}, database={'name': 'bar', 'schemas': None},
).once() ).once()
flexmock(module).should_receive('ensure_databases_found') flexmock(module).should_receive('ensure_databases_found')
@ -256,7 +256,9 @@ def test_run_restore_restores_each_database():
storage=flexmock(), storage=flexmock(),
hooks=flexmock(), hooks=flexmock(),
local_borg_version=flexmock(), local_borg_version=flexmock(),
restore_arguments=flexmock(repository='repo', archive='archive', databases=flexmock()), restore_arguments=flexmock(
repository='repo', archive='archive', databases=flexmock(), schemas=None
),
global_arguments=flexmock(dry_run=False), global_arguments=flexmock(dry_run=False),
local_path=flexmock(), local_path=flexmock(),
remote_path=flexmock(), remote_path=flexmock(),
@ -327,7 +329,7 @@ def test_run_restore_restores_database_configured_with_all_name():
remote_path=object, remote_path=object,
archive_name=object, archive_name=object,
hook_name='postgresql_databases', hook_name='postgresql_databases',
database={'name': 'foo'}, database={'name': 'foo', 'schemas': None},
).once() ).once()
flexmock(module).should_receive('restore_single_database').with_args( flexmock(module).should_receive('restore_single_database').with_args(
repository=object, repository=object,
@ -340,7 +342,7 @@ def test_run_restore_restores_database_configured_with_all_name():
remote_path=object, remote_path=object,
archive_name=object, archive_name=object,
hook_name='postgresql_databases', hook_name='postgresql_databases',
database={'name': 'bar'}, database={'name': 'bar', 'schemas': None},
).once() ).once()
flexmock(module).should_receive('ensure_databases_found') flexmock(module).should_receive('ensure_databases_found')
@ -350,7 +352,9 @@ def test_run_restore_restores_database_configured_with_all_name():
storage=flexmock(), storage=flexmock(),
hooks=flexmock(), hooks=flexmock(),
local_borg_version=flexmock(), local_borg_version=flexmock(),
restore_arguments=flexmock(repository='repo', archive='archive', databases=flexmock()), restore_arguments=flexmock(
repository='repo', archive='archive', databases=flexmock(), schemas=None
),
global_arguments=flexmock(dry_run=False), global_arguments=flexmock(dry_run=False),
local_path=flexmock(), local_path=flexmock(),
remote_path=flexmock(), remote_path=flexmock(),
@ -399,7 +403,7 @@ def test_run_restore_skips_missing_database():
remote_path=object, remote_path=object,
archive_name=object, archive_name=object,
hook_name='postgresql_databases', hook_name='postgresql_databases',
database={'name': 'foo'}, database={'name': 'foo', 'schemas': None},
).once() ).once()
flexmock(module).should_receive('restore_single_database').with_args( flexmock(module).should_receive('restore_single_database').with_args(
repository=object, repository=object,
@ -412,7 +416,7 @@ def test_run_restore_skips_missing_database():
remote_path=object, remote_path=object,
archive_name=object, archive_name=object,
hook_name='postgresql_databases', hook_name='postgresql_databases',
database={'name': 'bar'}, database={'name': 'bar', 'schemas': None},
).never() ).never()
flexmock(module).should_receive('ensure_databases_found') flexmock(module).should_receive('ensure_databases_found')
@ -422,7 +426,9 @@ def test_run_restore_skips_missing_database():
storage=flexmock(), storage=flexmock(),
hooks=flexmock(), hooks=flexmock(),
local_borg_version=flexmock(), local_borg_version=flexmock(),
restore_arguments=flexmock(repository='repo', archive='archive', databases=flexmock()), restore_arguments=flexmock(
repository='repo', archive='archive', databases=flexmock(), schemas=None
),
global_arguments=flexmock(dry_run=False), global_arguments=flexmock(dry_run=False),
local_path=flexmock(), local_path=flexmock(),
remote_path=flexmock(), remote_path=flexmock(),
@ -465,7 +471,7 @@ def test_run_restore_restores_databases_from_different_hooks():
remote_path=object, remote_path=object,
archive_name=object, archive_name=object,
hook_name='postgresql_databases', hook_name='postgresql_databases',
database={'name': 'foo'}, database={'name': 'foo', 'schemas': None},
).once() ).once()
flexmock(module).should_receive('restore_single_database').with_args( flexmock(module).should_receive('restore_single_database').with_args(
repository=object, repository=object,
@ -478,7 +484,7 @@ def test_run_restore_restores_databases_from_different_hooks():
remote_path=object, remote_path=object,
archive_name=object, archive_name=object,
hook_name='mysql_databases', hook_name='mysql_databases',
database={'name': 'bar'}, database={'name': 'bar', 'schemas': None},
).once() ).once()
flexmock(module).should_receive('ensure_databases_found') flexmock(module).should_receive('ensure_databases_found')
@ -488,7 +494,9 @@ def test_run_restore_restores_databases_from_different_hooks():
storage=flexmock(), storage=flexmock(),
hooks=flexmock(), hooks=flexmock(),
local_borg_version=flexmock(), local_borg_version=flexmock(),
restore_arguments=flexmock(repository='repo', archive='archive', databases=flexmock()), restore_arguments=flexmock(
repository='repo', archive='archive', databases=flexmock(), schemas=None
),
global_arguments=flexmock(dry_run=False), global_arguments=flexmock(dry_run=False),
local_path=flexmock(), local_path=flexmock(),
remote_path=flexmock(), remote_path=flexmock(),

View file

@ -157,7 +157,7 @@ def test_dump_databases_runs_mongodumpall_for_all_databases():
def test_restore_database_dump_runs_mongorestore(): def test_restore_database_dump_runs_mongorestore():
database_config = [{'name': 'foo'}] database_config = [{'name': 'foo', 'schemas': None}]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
flexmock(module).should_receive('make_dump_path') flexmock(module).should_receive('make_dump_path')
@ -189,7 +189,9 @@ def test_restore_database_dump_errors_on_multiple_database_config():
def test_restore_database_dump_runs_mongorestore_with_hostname_and_port(): def test_restore_database_dump_runs_mongorestore_with_hostname_and_port():
database_config = [{'name': 'foo', 'hostname': 'database.example.org', 'port': 5433}] database_config = [
{'name': 'foo', 'hostname': 'database.example.org', 'port': 5433, 'schemas': None}
]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
flexmock(module).should_receive('make_dump_path') flexmock(module).should_receive('make_dump_path')
@ -223,6 +225,7 @@ def test_restore_database_dump_runs_mongorestore_with_username_and_password():
'username': 'mongo', 'username': 'mongo',
'password': 'trustsome1', 'password': 'trustsome1',
'authentication_database': 'admin', 'authentication_database': 'admin',
'schemas': None,
} }
] ]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
@ -254,7 +257,7 @@ def test_restore_database_dump_runs_mongorestore_with_username_and_password():
def test_restore_database_dump_runs_mongorestore_with_options(): def test_restore_database_dump_runs_mongorestore_with_options():
database_config = [{'name': 'foo', 'restore_options': '--harder'}] database_config = [{'name': 'foo', 'restore_options': '--harder', 'schemas': None}]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
flexmock(module).should_receive('make_dump_path') flexmock(module).should_receive('make_dump_path')
@ -271,8 +274,36 @@ def test_restore_database_dump_runs_mongorestore_with_options():
) )
def test_restore_databases_dump_runs_mongorestore_with_schemas():
database_config = [{'name': 'foo', 'schemas': ['bar', 'baz']}]
extract_process = flexmock(stdout=flexmock())
flexmock(module).should_receive('make_dump_path')
flexmock(module.dump).should_receive('make_database_dump_filename')
flexmock(module).should_receive('execute_command_with_processes').with_args(
[
'mongorestore',
'--archive',
'--drop',
'--db',
'foo',
'--nsInclude',
'bar',
'--nsInclude',
'baz',
],
processes=[extract_process],
output_log_level=logging.DEBUG,
input_file=extract_process.stdout,
).once()
module.restore_database_dump(
database_config, 'test.yaml', {}, dry_run=False, extract_process=extract_process
)
def test_restore_database_dump_runs_psql_for_all_database_dump(): def test_restore_database_dump_runs_psql_for_all_database_dump():
database_config = [{'name': 'all'}] database_config = [{'name': 'all', 'schemas': None}]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
flexmock(module).should_receive('make_dump_path') flexmock(module).should_receive('make_dump_path')
@ -290,7 +321,7 @@ def test_restore_database_dump_runs_psql_for_all_database_dump():
def test_restore_database_dump_with_dry_run_skips_restore(): def test_restore_database_dump_with_dry_run_skips_restore():
database_config = [{'name': 'foo'}] database_config = [{'name': 'foo', 'schemas': None}]
flexmock(module).should_receive('make_dump_path') flexmock(module).should_receive('make_dump_path')
flexmock(module.dump).should_receive('make_database_dump_filename') flexmock(module.dump).should_receive('make_database_dump_filename')
@ -302,7 +333,7 @@ def test_restore_database_dump_with_dry_run_skips_restore():
def test_restore_database_dump_without_extract_process_restores_from_disk(): def test_restore_database_dump_without_extract_process_restores_from_disk():
database_config = [{'name': 'foo', 'format': 'directory'}] database_config = [{'name': 'foo', 'format': 'directory', 'schemas': None}]
flexmock(module).should_receive('make_dump_path') flexmock(module).should_receive('make_dump_path')
flexmock(module.dump).should_receive('make_database_dump_filename').and_return('/dump/path') flexmock(module.dump).should_receive('make_database_dump_filename').and_return('/dump/path')

View file

@ -411,7 +411,7 @@ def test_dump_databases_runs_non_default_pg_dump():
def test_restore_database_dump_runs_pg_restore(): def test_restore_database_dump_runs_pg_restore():
database_config = [{'name': 'foo'}] database_config = [{'name': 'foo', 'schemas': None}]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'}) flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
@ -458,7 +458,9 @@ def test_restore_database_dump_errors_on_multiple_database_config():
def test_restore_database_dump_runs_pg_restore_with_hostname_and_port(): def test_restore_database_dump_runs_pg_restore_with_hostname_and_port():
database_config = [{'name': 'foo', 'hostname': 'database.example.org', 'port': 5433}] database_config = [
{'name': 'foo', 'hostname': 'database.example.org', 'port': 5433, 'schemas': None}
]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'}) flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
@ -506,7 +508,9 @@ def test_restore_database_dump_runs_pg_restore_with_hostname_and_port():
def test_restore_database_dump_runs_pg_restore_with_username_and_password(): def test_restore_database_dump_runs_pg_restore_with_username_and_password():
database_config = [{'name': 'foo', 'username': 'postgres', 'password': 'trustsome1'}] database_config = [
{'name': 'foo', 'username': 'postgres', 'password': 'trustsome1', 'schemas': None}
]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
flexmock(module).should_receive('make_extra_environment').and_return( flexmock(module).should_receive('make_extra_environment').and_return(
@ -553,7 +557,12 @@ def test_restore_database_dump_runs_pg_restore_with_username_and_password():
def test_restore_database_dump_runs_pg_restore_with_options(): def test_restore_database_dump_runs_pg_restore_with_options():
database_config = [ database_config = [
{'name': 'foo', 'restore_options': '--harder', 'analyze_options': '--smarter'} {
'name': 'foo',
'restore_options': '--harder',
'analyze_options': '--smarter',
'schemas': None,
}
] ]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
@ -596,7 +605,7 @@ def test_restore_database_dump_runs_pg_restore_with_options():
def test_restore_database_dump_runs_psql_for_all_database_dump(): def test_restore_database_dump_runs_psql_for_all_database_dump():
database_config = [{'name': 'all'}] database_config = [{'name': 'all', 'schemas': None}]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'}) flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
@ -621,7 +630,12 @@ def test_restore_database_dump_runs_psql_for_all_database_dump():
def test_restore_database_dump_runs_non_default_pg_restore_and_psql(): def test_restore_database_dump_runs_non_default_pg_restore_and_psql():
database_config = [ database_config = [
{'name': 'foo', 'pg_restore_command': 'special_pg_restore', 'psql_command': 'special_psql'} {
'name': 'foo',
'pg_restore_command': 'special_pg_restore',
'psql_command': 'special_psql',
'schemas': None,
}
] ]
extract_process = flexmock(stdout=flexmock()) extract_process = flexmock(stdout=flexmock())
@ -654,7 +668,7 @@ def test_restore_database_dump_runs_non_default_pg_restore_and_psql():
def test_restore_database_dump_with_dry_run_skips_restore(): def test_restore_database_dump_with_dry_run_skips_restore():
database_config = [{'name': 'foo'}] database_config = [{'name': 'foo', 'schemas': None}]
flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'}) flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
flexmock(module).should_receive('make_dump_path') flexmock(module).should_receive('make_dump_path')
@ -667,7 +681,7 @@ def test_restore_database_dump_with_dry_run_skips_restore():
def test_restore_database_dump_without_extract_process_restores_from_disk(): def test_restore_database_dump_without_extract_process_restores_from_disk():
database_config = [{'name': 'foo'}] database_config = [{'name': 'foo', 'schemas': None}]
flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'}) flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
flexmock(module).should_receive('make_dump_path') flexmock(module).should_receive('make_dump_path')
@ -696,3 +710,39 @@ def test_restore_database_dump_without_extract_process_restores_from_disk():
module.restore_database_dump( module.restore_database_dump(
database_config, 'test.yaml', {}, dry_run=False, extract_process=None database_config, 'test.yaml', {}, dry_run=False, extract_process=None
) )
def test_restore_database_dump_with_schemas_restores_schemas():
database_config = [{'name': 'foo', 'schemas': ['bar', 'baz']}]
flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
flexmock(module).should_receive('make_dump_path')
flexmock(module.dump).should_receive('make_database_dump_filename').and_return('/dump/path')
flexmock(module).should_receive('execute_command_with_processes').with_args(
(
'pg_restore',
'--no-password',
'--if-exists',
'--exit-on-error',
'--clean',
'--dbname',
'foo',
'/dump/path',
'--schema',
'bar',
'--schema',
'baz',
),
processes=[],
output_log_level=logging.DEBUG,
input_file=None,
extra_environment={'PGSSLMODE': 'disable'},
).once()
flexmock(module).should_receive('execute_command').with_args(
('psql', '--no-password', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
extra_environment={'PGSSLMODE': 'disable'},
).once()
module.restore_database_dump(
database_config, 'test.yaml', {}, dry_run=False, extract_process=None
)