Add option to add specific center or doctor

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

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:

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