diff --git a/contrib/buildbot/phabricator_wrapper.py b/contrib/buildbot/phabricator_wrapper.py --- a/contrib/buildbot/phabricator_wrapper.py +++ b/contrib/buildbot/phabricator_wrapper.py @@ -484,3 +484,39 @@ buildTargetPHID=build_target.phid, type=harbormaster_build_status_mapping[build_target.status()] ) + + def get_object_token(self, object_PHID): + """ Return the current token set by the current user on target object """ + tokens = self.token.given( + authorPHIDs=[self.get_current_user_phid()], + objectPHIDs=[object_PHID], + tokenPHIDs=[], + ) + + if not tokens: + return "" + + # There should be no more than a single token from the same user for the + # same object. + if len(tokens) > 1: + self.logger.info( + "Found {} tokens for user {} on object {}: {}".format( + len(tokens), + self.get_current_user_phid(), + object_PHID, + tokens, + ) + ) + + return tokens[0]["tokenPHID"] + + def set_object_token(self, object_PHID, token_PHID=None): + """ Award or rescind a token for the target object """ + # If no token is given, rescind any previously awarded token + if token_PHID is None: + token_PHID = "" + + self.token.give( + objectPHID=object_PHID, + tokenPHID=token_PHID, + ) diff --git a/contrib/buildbot/test/mocks/phabricator.py b/contrib/buildbot/test/mocks/phabricator.py --- a/contrib/buildbot/test/mocks/phabricator.py +++ b/contrib/buildbot/test/mocks/phabricator.py @@ -54,6 +54,8 @@ phab.project = mock.Mock() phab.project.search.return_value = Result([]) + phab.token = mock.Mock() + phab.transaction = mock.Mock() phab.transaction.search.return_value = Result([]) diff --git a/contrib/buildbot/test/test_phabricator.py b/contrib/buildbot/test/test_phabricator.py --- a/contrib/buildbot/test/test_phabricator.py +++ b/contrib/buildbot/test/test_phabricator.py @@ -429,6 +429,76 @@ with self.assertRaisesRegex(AssertionError, "Failed to edit panel"): call_set_text_panel_content() + def test_get_object_token(self): + user_PHID = "PHID-USER-foobarbaz" + self.phab.user.whoami.return_value = { + "phid": user_PHID, + } + + object_PHID = "PHID-DREV-abcdef" + + def assert_token_given_called(): + self.phab.token.given.assert_called_with( + authorPHIDs=[user_PHID], + objectPHIDs=[object_PHID], + tokenPHIDs=[], + ) + + # There is no token for this object + self.phab.token.given.return_value = [] + token = self.phab.get_object_token(object_PHID) + assert_token_given_called() + self.assertEqual(token, "") + + # There is exactly 1 token for this object + self.phab.token.given.return_value = [ + { + "authorPHID": user_PHID, + "objectPHID": object_PHID, + "tokenPHID": "PHID-TOKN-like-1", + "dateCreated": 0, + }, + ] + token = self.phab.get_object_token(object_PHID) + assert_token_given_called() + self.assertEqual(token, "PHID-TOKN-like-1") + + # If there is more than a single token only the first one is returned + self.phab.token.given.return_value = [ + { + "authorPHID": user_PHID, + "objectPHID": object_PHID, + "tokenPHID": "PHID-TOKN-like-1", + "dateCreated": 0, + }, + { + "authorPHID": user_PHID, + "objectPHID": object_PHID, + "tokenPHID": "PHID-TOKN-like-2", + "dateCreated": 1, + }, + ] + token = self.phab.get_object_token(object_PHID) + assert_token_given_called() + self.assertEqual(token, "PHID-TOKN-like-1") + + def test_set_object_token(self): + object_PHID = "PHID-DREV-abcdef" + + def assert_token_give_called(token_PHID): + self.phab.token.give.assert_called_with( + objectPHID=object_PHID, + tokenPHID=token_PHID, + ) + + # Rescind any previoulsy awarded token + self.phab.set_object_token(object_PHID) + assert_token_give_called("") + + token_PHID = "PHID-TOKN-like-1" + self.phab.set_object_token(object_PHID, token_PHID) + assert_token_give_called(token_PHID) + if __name__ == '__main__': unittest.main()