diff --git a/miraheze/mediawiki/mwimport.py b/miraheze/mediawiki/mwimport.py index 9ed52b0..839f9f1 100644 --- a/miraheze/mediawiki/mwimport.py +++ b/miraheze/mediawiki/mwimport.py @@ -41,6 +41,9 @@ def parse_args(input_args: list | None = None, check_paths: bool = True) -> argp parser.add_argument('wiki', help='Database name of the wiki to import to') args = parser.parse_args(input_args) + if not args.wiki.endswith('wiki') and not args.wiki.endswith('wikibeta'): + raise ValueError(' must be a proper database name') + if not args.xml and not args.images: raise ValueError('--xml and/or --images must be passed') if args.images and not args.images_comment: diff --git a/tests/test_mwimport.py b/tests/test_mwimport.py index 6036938..f75913b 100644 --- a/tests/test_mwimport.py +++ b/tests/test_mwimport.py @@ -52,6 +52,22 @@ def test_parse_args_username_prefix(): assert args.username_prefix == 'w' +def test_parse_args_beta_database_name(): + args = mwimport.parse_args([ + '--xml=dump.xml', + 'examplewikibeta', + ], False) + assert args.wiki == 'examplewikibeta' + + +def test_parse_args_invalid_database_name(): + with pytest.raises(ValueError, match=' must be a proper database name'): + mwimport.parse_args([ + '--xml=dump.xml', + 'example', + ]) + + def test_parse_args_need_xml_or_images(): with pytest.raises(ValueError, match='--xml and/or --images must be passed'): mwimport.parse_args([