diff --git a/README.md b/README.md index f72631d..7f79c43 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,8 @@ Further optional arguments: exclude centers --center-exclude-regex CENTER_EXCLUDE_REGEX exclude centers by regex +--additional-center "url", -ac "url" + Add additional centers or doctors e.g. "/institut/berlin/ciz-berlin-berlin" --include-neighbor-city, -n include neighboring cities --start-date START_DATE diff --git a/doctoshotgun.py b/doctoshotgun.py index 8e9096e..6ebbb4c 100755 --- a/doctoshotgun.py +++ b/doctoshotgun.py @@ -617,6 +617,16 @@ def setup_loggers(self, level): logging.root.setLevel(level) logging.root.addHandler(self.create_default_logger()) + def try_to_book_or_sleep(self, docto, center, vaccine_list, start_date, end_date, only_second, only_third, dry_run=False): + if docto.try_to_book(center, vaccine_list, start_date, end_date, only_second, only_third, dry_run): + log('') + log('💉 %s Congratulations.' % + colored('Booked!', 'green', attrs=('bold',))) + return True + + sleep(SLEEP_INTERVAL_AFTER_CENTER) + return False + def main(self, cli_args=None): colorama.init() # needed for windows @@ -653,6 +663,8 @@ def main(self, cli_args=None): action='append', help='exclude centers') parser.add_argument('--center-exclude-regex', action='append', help='exclude centers by regex') + parser.add_argument('--additional-center', '-ac', + action='append', help='Add additional centers or doctors: "url" e.g. "/institut/berlin/ciz-berlin-berlin"') parser.add_argument( '--include-neighbor-city', '-n', action='store_true', help='include neighboring cities') parser.add_argument('--start-date', type=str, default=None, @@ -669,8 +681,6 @@ def main(self, cli_args=None): parser.add_argument('--code', type=str, default=None, help='2FA code') args = parser.parse_args(cli_args if cli_args else sys.argv[1:]) - from types import SimpleNamespace - if args.debug: responses_dirname = tempfile.mkdtemp(prefix='woob_session_') self.setup_loggers(logging.DEBUG) @@ -827,18 +837,19 @@ def main(self, cli_args=None): continue log('') - log('Center %(name_with_title)s (%(city)s):' % center) - if docto.try_to_book(center, vaccine_list, start_date, end_date, args.only_second, args.only_third, args.dry_run): - log('') - log('💉 %s Congratulations.' % - colored('Booked!', 'green', attrs=('bold',))) + if self.try_to_book_or_sleep(docto, center, vaccine_list, start_date, end_date, args.only_second, args.only_third, args.dry_run): return 0 - sleep(SLEEP_INTERVAL_AFTER_CENTER) + if args.additional_center: + for additional_center in args.additional_center: + log('') + log('Additional center %s:', additional_center.split('/')[-1]) + + if self.try_to_book_or_sleep(docto, { "url": additional_center }, vaccine_list, start_date, end_date, args.only_second, args.only_third, args.dry_run): + return 0 - log('') log('No free slots found at selected centers. Trying another round in %s sec...', SLEEP_INTERVAL_AFTER_RUN) sleep(SLEEP_INTERVAL_AFTER_RUN) except CityNotFound as e: diff --git a/test_cli_args.py b/test_cli_args.py index 3fcb562..20b4d6e 100644 --- a/test_cli_args.py +++ b/test_cli_args.py @@ -63,6 +63,36 @@ def test_center_exclude_arg_should_filter_excluded_centers(MockDoctolibDE, tmp_p assert call_args_list.args[0]['city'] == city +@responses.activate +@patch('doctoshotgun.DoctolibDE') +def test_additional_centers_arg_should_processed(MockDoctolibDE, tmp_path): + """ + Check that additional centers are processed + """ + # prepare + mock_doctolib_de = get_mocked_doctolib(MockDoctolibDE) + + # patch find_centers to force usage of additional centers + mock_doctolib_de.find_centers.return_value = {} + + # call for every additional center as every booking is successfull + city = '' + additional_centers = [ + '/institut/koeln/ciz-koeln-koeln', + '/allgemeinmedizin/koeln/dr-dre' + ] + + for additional_center in additional_centers: + call_application(city, cli_args=['--additional-center', additional_center]) + + # assert + assert mock_doctolib_de.get_patients.called + assert mock_doctolib_de.try_to_book.called + + for i, call_args_list in enumerate(mock_doctolib_de.try_to_book.call_args_list): + assert call_args_list[0][0]['url'] == additional_centers[i] + + def get_mocked_doctolib(MockDoctolibDE): mock_doctolib_de = MagicMock(wraps=DoctolibDE) MockDoctolibDE.return_value = mock_doctolib_de