From c4d6bb38a15dee0f896948c6cb9de385c1ddc6d9 Mon Sep 17 00:00:00 2001
From: Alan Hamlett <alan.hamlett@gmail.com>
Date: Wed, 31 Aug 2016 14:27:10 +0200
Subject: [PATCH] test for unknown exceptions and malformed extra heartbeats
 json

---
 tests/test_main.py    | 62 +++++++++++++++++++++++++++++++++++++++++++
 wakatime/constants.py | 24 ++++++++++++++++-
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/tests/test_main.py b/tests/test_main.py
index 5b5e8e0..9bd29b5 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -17,6 +17,7 @@ from wakatime.constants import (
     AUTH_ERROR,
     CONFIG_FILE_PARSE_ERROR,
     SUCCESS,
+    MALFORMED_HEARTBEAT_ERROR,
 )
 from wakatime.packages.requests.models import Response
 from . import utils
@@ -718,6 +719,43 @@ class MainTestCase(utils.TestCase):
                 self.assertEquals(data.get('entity'), entity2)
                 self.assertEquals(data.get('project'), project2)
 
+    @log_capture()
+    def test_extra_heartbeats_with_malformed_json(self, logs):
+        logging.disable(logging.NOTSET)
+
+        response = Response()
+        response.status_code = 201
+        self.patched['wakatime.packages.requests.adapters.HTTPAdapter.send'].return_value = response
+
+        with utils.TemporaryDirectory() as tempdir:
+            entity = 'tests/samples/codefiles/twolinefile.txt'
+            shutil.copy(entity, os.path.join(tempdir, 'twolinefile.txt'))
+            entity = os.path.realpath(os.path.join(tempdir, 'twolinefile.txt'))
+
+            entity = os.path.abspath('tests/samples/codefiles/emptyfile.txt')
+            config = 'tests/samples/configs/good_config.cfg'
+            args = ['--file', entity, '--config', config, '--extra-heartbeats']
+
+            with utils.mock.patch('wakatime.main.sys.stdin') as mock_stdin:
+                heartbeats = '[{foobar}]'
+                mock_stdin.readline.return_value = heartbeats
+
+                retval = execute(args)
+
+                self.assertEquals(retval, MALFORMED_HEARTBEAT_ERROR)
+                self.assertEquals(sys.stdout.getvalue(), '')
+                self.assertEquals(sys.stderr.getvalue(), '')
+
+                log_output = u("\n").join([u(' ').join(x) for x in logs.actual()])
+                self.assertEquals(log_output, '')
+
+                self.patched['wakatime.session_cache.SessionCache.get'].assert_called_once_with()
+                self.patched['wakatime.session_cache.SessionCache.delete'].assert_not_called()
+                self.patched['wakatime.session_cache.SessionCache.save'].assert_called_once_with(ANY)
+
+                self.patched['wakatime.offlinequeue.Queue.push'].assert_not_called()
+                self.patched['wakatime.offlinequeue.Queue.pop'].assert_not_called()
+
     @log_capture()
     def test_nonascii_filename(self, logs):
         logging.disable(logging.NOTSET)
@@ -770,3 +808,27 @@ class MainTestCase(utils.TestCase):
                 self.assertEquals(heartbeat[key], val)
             self.assertEquals(stats, json.loads(self.patched['wakatime.offlinequeue.Queue.push'].call_args[0][1]))
             self.patched['wakatime.offlinequeue.Queue.pop'].assert_not_called()
+
+    @log_capture()
+    def test_unhandled_exception(self, logs):
+        logging.disable(logging.NOTSET)
+
+        with utils.mock.patch('wakatime.main.process_heartbeat') as mock_process_heartbeat:
+            ex_msg = 'testing unhandled exception'
+            mock_process_heartbeat.side_effect = RuntimeError(ex_msg)
+
+            entity = 'tests/samples/codefiles/twolinefile.txt'
+            config = 'tests/samples/configs/good_config.cfg'
+            args = ['--entity', entity, '--key', '123', '--config', config]
+
+            execute(args)
+
+            self.assertIn(ex_msg, sys.stdout.getvalue())
+            self.assertEquals(sys.stderr.getvalue(), '')
+
+            log_output = u("\n").join([u(' ').join(x) for x in logs.actual()])
+            self.assertIn(ex_msg, log_output)
+
+            self.patched['wakatime.offlinequeue.Queue.push'].assert_not_called()
+            self.patched['wakatime.offlinequeue.Queue.pop'].assert_not_called()
+            self.patched['wakatime.session_cache.SessionCache.get'].assert_not_called()
diff --git a/wakatime/constants.py b/wakatime/constants.py
index 1c7316b..fece2b9 100644
--- a/wakatime/constants.py
+++ b/wakatime/constants.py
@@ -9,10 +9,32 @@
     :license: BSD, see LICENSE for more details.
 """
 
-
+""" Success
+Exit code used when a heartbeat was sent successfully.
+"""
 SUCCESS = 0
+
+""" Api Error
+Exit code used when the WakaTime API returned an error.
+"""
 API_ERROR = 102
+
+""" Config File Parse Error
+Exit code used when the ~/.wakatime.cfg config file could not be parsed.
+"""
 CONFIG_FILE_PARSE_ERROR = 103
+
+""" Auth Error
+Exit code used when our api key is invalid.
+"""
 AUTH_ERROR = 104
+
+""" Unknown Error
+Exit code used when there was an unhandled exception.
+"""
 UNKNOWN_ERROR = 105
+
+""" Malformed Heartbeat Error
+Exit code used when the JSON input from `--extra-heartbeats` is malformed.
+"""
 MALFORMED_HEARTBEAT_ERROR = 106