diff --git a/jobrunner/mail/chat.py b/jobrunner/mail/chat.py index 07c6a29..b90bb00 100644 --- a/jobrunner/mail/chat.py +++ b/jobrunner/mail/chat.py @@ -77,7 +77,7 @@ class PostError(Exception): pass -def _postToGChat(text, uri, threadId=None): +def _postToGChat(text, uri, timeout, threadId=None): payload = { "text": text, } @@ -86,7 +86,7 @@ def _postToGChat(text, uri, threadId=None): } if threadId: payload["thread"] = {"name": threadId} - ret = requests.post(uri, json=payload, headers=headers) + ret = requests.post(uri, json=payload, headers=headers, timeout=timeout) try: return ret.json()["thread"]["name"] except (KeyError, TypeError): @@ -154,6 +154,9 @@ def parseArgs(args=None): ap.add_argument("-f", dest="inFile", action="store") ap.add_argument("-T", "--new-thread", action="store_true", help="Do not re-use the previous chat thread for each hook") + ap.add_argument("--timeout", type=int, default=10, + help="Specify timeout (in seconds) for REST API requests " + "(default=%(default)s)") ap.add_argument("toAddr", metavar="to-addr", nargs="+") return ap.parse_args(args) @@ -191,7 +194,7 @@ def main(args=None): if not opts.new_thread and config.chatmailReuseThreads else None) - newThreadId = _postToGChat(msg, hook, threadId=threadId) + newThreadId = _postToGChat(msg, hook, opts.timeout, threadId=threadId) threadCache.put(hook, newThreadId) return OK diff --git a/jobrunner/test/chatmail_test.py b/jobrunner/test/chatmail_test.py index e952ebc..7ebff29 100644 --- a/jobrunner/test/chatmail_test.py +++ b/jobrunner/test/chatmail_test.py @@ -13,6 +13,8 @@ # pylint: disable-msg=protected-access +TIMEOUT = 10 + @patch("jobrunner.mail.chat.requests", new=MagicMock(requests)) class ChatPostTest(TestCase): @@ -21,7 +23,7 @@ def testGoodRequest(self): chat.requests.post = MagicMock(return_value=retMock) retMock.json.return_value = {"thread": {"name": "somethread"}} self.assertEqual( - chat._postToGChat("foo", "https://something"), + chat._postToGChat("foo", "https://something", TIMEOUT), "somethread") def testBadRequestReturn(self): @@ -36,7 +38,7 @@ def testBadRequestReturn(self): for retVal in retValues: retMock.json.return_value = retVal self.assertEqual( - chat._postToGChat("foo", "https://something"), + chat._postToGChat("foo", "https://something", TIMEOUT), None) @@ -93,7 +95,7 @@ def assertPostTextMatchesRegexp(self, regexp, callIdx=0): def testRequestSafetySanity(self, *_): with self.assertRaises(AttributeError): - chat.requests.post() + chat.requests.post("http://example.com/foo") def testBasicCall(self, *mockArgs): mocks = ChatTest.Mocks(mockArgs) @@ -102,7 +104,7 @@ def testBasicCall(self, *mockArgs): chat.main(["user1"]) - mocks.postToGChat.assert_called_once_with(ANY, hook, threadId=None) + mocks.postToGChat.assert_called_once_with(ANY, hook, TIMEOUT, threadId=None) self.assertPostTextMatchesRegexp("NO SUBJECT") def testBasicCallWithOptions(self, *mockArgs): @@ -123,7 +125,7 @@ def testBasicCallWithOptions(self, *mockArgs): openMock.assert_called_once_with("mytext.txt") self.assertEqual(ret, 0) - mocks.postToGChat.assert_called_once_with(ANY, hook, threadId=None) + mocks.postToGChat.assert_called_once_with(ANY, hook, TIMEOUT, threadId=None) self.assertPostTextMatchesRegexp( r"My subject.*a bunch of\ntext.*attachfile.txt") @@ -140,7 +142,8 @@ def testPipeContent(self, *mockArgs): ret = chat.main(["-s", "My subject", "user1"]) self.assertEqual(ret, 0) - mocks.postToGChat.assert_called_once_with(ANY, hook, threadId=None) + mocks.postToGChat.assert_called_once_with(ANY, hook, TIMEOUT, + threadId=None) self.assertPostTextMatchesRegexp( r"\*My subject\*" + (r"\s*$" if tty else r".*a bunch of\ntext")) @@ -159,7 +162,7 @@ def testCallWithMultipleHooks(self, *mockArgs): "user1", "user2"]) self.assertEqual(ret, 0) - calls = [call(ANY, hook, threadId=None) + calls = [call(ANY, hook, TIMEOUT, threadId=None) for hook in set(hooks.values())] self.assertEqual(mocks.postToGChat.call_count, len(calls)) mocks.postToGChat.assert_has_calls(calls, any_order=True) @@ -196,7 +199,7 @@ def _runTest(retThreadId, reuseThreads): self.assertEqual(ret, 0) mocks.postToGChat.assert_called_once_with( - ANY, hook, + ANY, hook, TIMEOUT, threadId=threadId if reuseThreads else None) if threadId == retThreadId: self.assertEqual(mocks.threadCacheWrite.call_count, 0) @@ -224,7 +227,8 @@ def _runTest(atAllConfig, userIdDict, expectAtAll=False): ret = chat.main(["-s", "My subject"] + users) self.assertEqual(ret, 0) - mocks.postToGChat.assert_called_once_with(ANY, hook, threadId=None) + mocks.postToGChat.assert_called_once_with(ANY, hook, TIMEOUT, + threadId=None) pattern = r"\*My subject\*" tags = [] if expectAtAll: