Update quick to prompt for cal when ambiguous, like add already does

Fixes #482.

gcalcli/gcal.py

@@ -1267,20 +1267,37 @@ def CalQuery(self, cmd, start_text='', count=1):
 
         self._GraphEvents(cmd, start, count, event_list)
 
+    def _prompt_for_calendar(self, cals):
+        if not cals:
+            raise GcalcliError('No avaliable calendar to use')
+        if len(cals) == 1:
+            return cals[0]
+        # Calendar not specified. Prompt the user to select it
+        writers = (self.ACCESS_OWNER, self.ACCESS_WRITER)
+        cals_with_write_perms = [cal for cal in self.cals
+                                 if cal['accessRole'] in writers]
+
+        print('\n'.join(f"{idx} {cal['summary']}"
+                        for idx, cal in enumerate(cals_with_write_perms)))
+        val = get_input(self.printer, 'Specify calendar from above: ',
+                        STR_TO_INT)
+        try:
+            return cals_with_write_perms[int(val)]
+        except IndexError:
+            raise GcalcliError(
+                f'Invalid selection from the list above: {val}\n')
+
     def QuickAddEvent(self, event_text, reminders=None):
         """Wrapper around Google Calendar API's quickAdd"""
         if not event_text:
             raise GcalcliError('event_text is required for a quickAdd')
 
-        if len(self.cals) != 1:
-            # TODO: get a better name for this exception class
-            # and use it elsewhere
-            raise GcalcliError('You must only specify a single calendar\n')
+        calendar = self._prompt_for_calendar(self.cals)
 
         new_event = self._retry_with_backoff(
             self.get_events()
                 .quickAdd(
-                    calendarId=self.cals[0]['id'],
+                    calendarId=calendar['id'],
                     text=event_text
                 )
         )
@@ -1297,7 +1314,7 @@ def QuickAddEvent(self, event_text, reminders=None):
             new_event = self._retry_with_backoff(
                             self.get_events()
                                 .patch(
-                                    calendarId=self.cals[0]['id'],
+                                    calendarId=calendar['id'],
                                     eventId=new_event['id'],
                                     body=rem
                                 )
@@ -1311,24 +1328,7 @@ def QuickAddEvent(self, event_text, reminders=None):
 
     def AddEvent(self, title, where, start, end, descr, who, reminders, color):
 
-        if len(self.cals) != 1:
-            # Calendar not specified. Prompt the user to select it
-            writers = (self.ACCESS_OWNER, self.ACCESS_WRITER)
-            cals_with_write_perms = [cal for cal in self.cals
-                                     if cal['accessRole'] in writers]
-
-            cal_names_with_idx = []
-            for idx, cal in enumerate(cals_with_write_perms):
-                cal_names_with_idx.append(str(idx) + ' ' + cal['summary'])
-            cal_names_with_idx = '\n'.join(cal_names_with_idx)
-            print(cal_names_with_idx)
-            val = get_input(self.printer, 'Specify calendar from above: ',
-                            STR_TO_INT)
-            try:
-                self.cals = [cals_with_write_perms[int(val)]]
-            except IndexError:
-                raise GcalcliError("The entered number doesn't appear on the "
-                                   'list above\n')
+        calendar = self._prompt_for_calendar(self.cals)
 
         event = {}
         event['summary'] = title
@@ -1339,9 +1339,9 @@ def AddEvent(self, title, where, start, end, descr, who, reminders, color):
 
         else:
             event['start'] = {'dateTime': start,
-                              'timeZone': self.cals[0]['timeZone']}
+                              'timeZone': calendar['timeZone']}
             event['end'] = {'dateTime': end,
-                            'timeZone': self.cals[0]['timeZone']}
+                            'timeZone': calendar['timeZone']}
 
         if where:
             event['location'] = where
@@ -1355,7 +1355,7 @@ def AddEvent(self, title, where, start, end, descr, who, reminders, color):
 
         event = self._add_reminders(event, reminders)
         events = self.get_events()
-        request = events.insert(calendarId=self.cals[0]['id'], body=event)
+        request = events.insert(calendarId=calendar['id'], body=event)
         new_event = self._retry_with_backoff(request)
 
         if self.details.get('url'):

tests/test_gcalcli.py

@@ -1,8 +1,10 @@
 from __future__ import absolute_import
 
+import io
 import os
 from datetime import datetime
 from json import load
+import re
 
 import pytest
 from dateutil.tz import tzutc
@@ -100,16 +102,34 @@ def test_add_event(PatchedGCalI):
     cal_names = parse_cal_names(['[email protected]'])
     gcal = PatchedGCalI(
             cal_names=cal_names, allday=False, default_reminders=True)
-    title = 'test event'
-    where = 'anywhere'
-    start = 'now'
-    end = 'tomorrow'
-    descr = 'testing'
-    who = 'anyone'
-    reminders = None
-    color = "banana"
-    assert gcal.AddEvent(
-        title, where, start, end, descr, who, reminders, color)
+    assert gcal.AddEvent(title='test event',
+                         where='anywhere',
+                         start='now',
+                         end='tomorrow',
+                         descr='testing',
+                         who='anyone',
+                         reminders=None,
+                         color='banana')
+
+
+def test_add_event_with_cal_prompt(PatchedGCalI, capsys, monkeypatch):
+    cal_names = parse_cal_names(['[email protected]', '[email protected]'])
+    gcal = PatchedGCalI(
+            cal_names=cal_names, allday=False, default_reminders=True)
+    # Fake selecting calendar 0 at the prompt
+    monkeypatch.setattr('sys.stdin', io.StringIO('0\n'))
+    assert gcal.AddEvent(title='test event',
+                         where='',
+                         start='now',
+                         end='tomorrow',
+                         descr='',
+                         who='',
+                         reminders=None,
+                         color='')
+    captured = capsys.readouterr()
+    assert re.match(
+        r'(?sm)^0 .*\n1 .*\n.*Specify calendar.*$', captured.out), \
+        f'Unexpected stderr: {captured.out}'
 
 
 def test_add_event_override_color(capsys, default_options,
@@ -127,9 +147,23 @@ def test_add_event_override_color(capsys, default_options,
 def test_quick_add(PatchedGCalI):
     cal_names = parse_cal_names(['[email protected]'])
     gcal = PatchedGCalI(cal_names=cal_names)
-    event_text = 'quick test event'
-    reminder = '5m sms'
-    assert gcal.QuickAddEvent(event_text, reminders=[reminder])
+    assert gcal.QuickAddEvent(
+        event_text='quick test event',
+        reminders=['5m sms'])
+
+
+def test_quick_add_with_cal_prompt(PatchedGCalI, capsys, monkeypatch):
+    cal_names = parse_cal_names(['[email protected]', '[email protected]'])
+    gcal = PatchedGCalI(cal_names=cal_names)
+    # Fake selecting calendar 0 at the prompt
+    monkeypatch.setattr('sys.stdin', io.StringIO('0\n'))
+    assert gcal.QuickAddEvent(
+        event_text='quick test event',
+        reminders=['5m sms'])
+    captured = capsys.readouterr()
+    assert re.match(
+        r'(?sm)^0 .*\n1 .*\n.*Specify calendar.*$', captured.out), \
+        f'Unexpected stderr: {captured.out}'
 
 
 def test_text_query(PatchedGCalI):