rearranged prompt_for_lab_award a bit to be more clear; removed unnee…

…ded test
4dn-dcic · Apr 22, 2019 · 1accb04 · 1accb04
1 parent 9ebc709
commit 1accb04
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 31 deletions.
diff --git a/tests/test_get_field_info.py b/tests/test_get_field_info.py
@@ -185,21 +185,6 @@ def test_connection_prompt_for_lab_award_multi_lab_award(
         assert connection.award == chosenaward
 
 
-def test_connection_no_prompt_if_no_labs(mocker, mkey, returned_user_me_submit_for_one_lab,
-                                         returned_lab_w_one_award):
-    '''using as remote user who has only one lab submits_for that has one award
-       in this case checking to make sure that lab and award are still present on connection
-    '''
-    with mocker.patch('dcicutils.ff_utils.get_metadata',
-                      side_effect=[returned_user_me_submit_for_one_lab.json(),
-                                   returned_lab_w_one_award.json()]):
-        connection = gfi.FDN_Connection(mkey)
-        connection.labs = None
-        connection.prompt_for_lab_award()
-        assert connection.labs is None
-        assert connection.award is None
-
-
 def test_set_award_no_lab(mocker, mkey, returned_user_me_submit_for_one_lab,
                           returned_lab_w_one_award):
     with mocker.patch('dcicutils.ff_utils.get_metadata',

diff --git a/wranglertools/get_field_info.py b/wranglertools/get_field_info.py
@@ -182,25 +182,26 @@ def prompt_for_lab_award(self, lab=None, award=None):
             has multiple awards and if so prompts for the one to set
             for the connection
         '''
-        if not lab and self.labs is not None:
-            if len(self.labs) > 1:
-                lchoices = []
-                print("Submitting for multiple labs:")
-                for i, lab in enumerate(self.labs):
-                    ch = str(i + 1)
-                    lchoices.append(ch)
-                    print("  ({choice}) {labname}".format(choice=ch, labname=lab))
-                lab_resp = str(input("Select the lab for this connection {choices}: ".format(choices=lchoices)))
-                if lab_resp not in lchoices:
-                    print("Not a valid choice - using {default}".format(default=self.lab))
-                    return
-                else:
-                    self.lab = self.labs[int(lab_resp) - 1]
+        if lab:
             if not award:
-                self.set_award(self.lab, False)
+                self.set_award(self.lab)
         else:
+            if self.labs is not None:
+                if len(self.labs) > 1:
+                    lchoices = []
+                    print("Submitting for multiple labs:")
+                    for i, lab in enumerate(self.labs):
+                        ch = str(i + 1)
+                        lchoices.append(ch)
+                        print("  ({choice}) {labname}".format(choice=ch, labname=lab))
+                    lab_resp = str(input("Select the lab for this connection {choices}: ".format(choices=lchoices)))
+                    if lab_resp not in lchoices:
+                        print("Not a valid choice - using {default}".format(default=self.lab))
+                        return
+                    else:
+                        self.lab = self.labs[int(lab_resp) - 1]
             if not award:
-                self.set_award(None)
+                self.set_award(self.lab, False)
 
 
 @attr.s