new --extra-heartbeats argument

2016-04-21 00:30:44 +02:00 · 2016-04-21 00:30:44 +02:00 · 67a2403070
parent 984a8ab62a
commit 67a2403070
5 changed files with 107 additions and 45 deletions
--- a/tests/samples/output/main_test_timeout_passed_via_command_line
+++ b/tests/samples/output/main_test_timeout_passed_via_command_line
@ -4,7 +4,7 @@ usage: wakatime [-h] [--entity FILE] [--key KEY] [--write] [--plugin PLUGIN]
                [--alternate-project ALTERNATE_PROJECT]
                [--alternate-language ALTERNATE_LANGUAGE]
                [--hostname HOSTNAME] [--disableoffline] [--hidefilenames]
-                [--exclude EXCLUDE] [--include INCLUDE] [--logfile LOGFILE]
-                [--apiurl API_URL] [--timeout TIMEOUT] [--config CONFIG]
-                [--verbose] [--version]
+                [--exclude EXCLUDE] [--include INCLUDE] [--extra-heartbeats]
+                [--logfile LOGFILE] [--apiurl API_URL] [--timeout TIMEOUT]
+                [--config CONFIG] [--verbose] [--version]
 wakatime: error: argument --timeout: invalid int value: 'abc'
--- a/tests/samples/output/test_help_contents
+++ b/tests/samples/output/test_help_contents
@ -4,9 +4,9 @@ usage: wakatime [-h] [--entity FILE] [--key KEY] [--write] [--plugin PLUGIN]
                [--alternate-project ALTERNATE_PROJECT]
                [--alternate-language ALTERNATE_LANGUAGE]
                [--hostname HOSTNAME] [--disableoffline] [--hidefilenames]
-                [--exclude EXCLUDE] [--include INCLUDE] [--logfile LOGFILE]
-                [--apiurl API_URL] [--timeout TIMEOUT] [--config CONFIG]
-                [--verbose] [--version]
+                [--exclude EXCLUDE] [--include INCLUDE] [--extra-heartbeats]
+                [--logfile LOGFILE] [--apiurl API_URL] [--timeout TIMEOUT]
+                [--config CONFIG] [--verbose] [--version]

 Common interface for the WakaTime api.

@ -46,6 +46,8 @@ optional arguments:
  --include INCLUDE     filename patterns to log; when used in combination
                        with --exclude, files matching include will still be
                        logged; POSIX regex syntax; can be used more than once
+  --extra-heartbeats    reads extra heartbeats from STDIN as a JSON array
+                        until EOF
  --logfile LOGFILE     defaults to ~/.wakatime.log
  --apiurl API_URL      heartbeats api url; for debugging with a local server
  --timeout TIMEOUT     number of seconds to wait when sending heartbeats to
--- a/tests/samples/output/test_missing_config_file
+++ b/tests/samples/output/test_missing_config_file
@ -4,7 +4,7 @@ usage: wakatime [-h] [--entity FILE] [--key KEY] [--write] [--plugin PLUGIN]
                [--alternate-project ALTERNATE_PROJECT]
                [--alternate-language ALTERNATE_LANGUAGE]
                [--hostname HOSTNAME] [--disableoffline] [--hidefilenames]
-                [--exclude EXCLUDE] [--include INCLUDE] [--logfile LOGFILE]
-                [--apiurl API_URL] [--timeout TIMEOUT] [--config CONFIG]
-                [--verbose] [--version]
+                [--exclude EXCLUDE] [--include INCLUDE] [--extra-heartbeats]
+                [--logfile LOGFILE] [--apiurl API_URL] [--timeout TIMEOUT]
+                [--config CONFIG] [--verbose] [--version]
 wakatime: error: Missing api key
--- a/tests/test_main.py
+++ b/tests/test_main.py
@ -23,9 +23,9 @@ try:
 except (ImportError, SyntaxError):
    import json
 try:
-    from mock import ANY
+    from mock import ANY, call
 except ImportError:
-    from unittest.mock import ANY
+    from unittest.mock import ANY, call
 from wakatime.packages import tzlocal


@ -647,3 +647,35 @@ class BaseTestCase(utils.TestCase):
        timezone = tzlocal.get_localzone()
        headers = self.patched['wakatime.packages.requests.adapters.HTTPAdapter.send'].call_args[0][0].headers
        self.assertEquals(headers.get('TimeZone'), u(timezone.zone).encode('utf-8') if is_py3 else timezone.zone)
+
+    def test_extra_heartbeats_argument(self):
+        response = Response()
+        response.status_code = 201
+        self.patched['wakatime.packages.requests.adapters.HTTPAdapter.send'].return_value = response
+
+        entity = '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:
+            now = int(time.time())
+            heartbeats = json.dumps([{
+                'timestamp': now,
+                'entity': entity,
+                'entity_type': 'file',
+                'is_write': True,
+            }])
+            mock_stdin.read.return_value = heartbeats
+
+            retval = execute(args)
+
+            self.assertEquals(retval, SUCCESS)
+            self.assertEquals(sys.stdout.getvalue(), '')
+            self.assertEquals(sys.stderr.getvalue(), '')
+
+            self.patched['wakatime.session_cache.SessionCache.get'].assert_has_calls([call(), call()])
+            self.patched['wakatime.session_cache.SessionCache.delete'].assert_not_called()
+            self.patched['wakatime.session_cache.SessionCache.save'].assert_has_calls([call(ANY), call(ANY)])
+
+            self.patched['wakatime.offlinequeue.Queue.push'].assert_not_called()
+            self.patched['wakatime.offlinequeue.Queue.pop'].assert_has_calls([call(), call()])
--- a/wakatime/main.py
+++ b/wakatime/main.py
@ -153,6 +153,9 @@ def parseArguments():
                 'POSIX regex syntax; can be used more than once')
    parser.add_argument('--ignore', dest='ignore', action='append',
            help=argparse.SUPPRESS)
+    parser.add_argument('--extra-heartbeats', dest='extra_heartbeats',
+            action='store_true',
+            help='reads extra heartbeats from STDIN as a JSON array until EOF')
    parser.add_argument('--logfile', dest='logfile',
            help='defaults to ~/.wakatime.log')
    parser.add_argument('--apiurl', dest='api_url',
@ -452,6 +455,47 @@ def sync_offline_heartbeats(args, hostname):
    return SUCCESS


+def process_heartbeat(args, configs, heartbeat):
+    exclude = should_exclude(args.entity, args.include, args.exclude)
+    if exclude is not False:
+        log.debug(u('Skipping because matches exclude pattern: {pattern}').format(
+            pattern=u(exclude),
+        ))
+        return SUCCESS
+
+    if args.entity_type != 'file' or os.path.isfile(args.entity):
+
+        stats = get_file_stats(args.entity,
+                                entity_type=args.entity_type,
+                                lineno=args.lineno,
+                                cursorpos=args.cursorpos,
+                                plugin=args.plugin,
+                                alternate_language=args.alternate_language)
+
+        project = args.project or args.alternate_project
+        branch = None
+        if args.entity_type == 'file':
+            project, branch = get_project_info(configs, args)
+
+        kwargs = vars(args)
+        kwargs['project'] = project
+        kwargs['branch'] = branch
+        kwargs['stats'] = stats
+        hostname = args.hostname or socket.gethostname()
+        kwargs['hostname'] = hostname
+        kwargs['timeout'] = args.timeout
+
+        status = send_heartbeat(**kwargs)
+        if status == SUCCESS:
+            return sync_offline_heartbeats(args, hostname)
+        else:
+            return status
+
+    else:
+        log.debug('File does not exist; ignoring this heartbeat.')
+        return SUCCESS
+
+
 def execute(argv=None):
    if argv:
        sys.argv = ['wakatime'] + argv
@ -463,44 +507,28 @@ def execute(argv=None):
    setup_logging(args, __version__)

    try:
-        exclude = should_exclude(args.entity, args.include, args.exclude)
-        if exclude is not False:
-            log.debug(u('Skipping because matches exclude pattern: {pattern}').format(
-                pattern=u(exclude),
-            ))
-            return SUCCESS

-        if args.entity_type != 'file' or os.path.isfile(args.entity):
+        heartbeat = {
+            'timestamp': args.timestamp,
+            'entity': args.entity,
+            'entity_type': args.entity_type,
+            'project': args.project,
+            'is_write': args.isWrite,
+            'lineno': args.lineno,
+            'cursorpos': args.cursorpos,
+            'alternate_project': args.alternate_project,
+            'alternate_language': args.alternate_language,
+        }
+        retval = process_heartbeat(args, configs, heartbeat)

-            stats = get_file_stats(args.entity,
-                                   entity_type=args.entity_type,
-                                   lineno=args.lineno,
-                                   cursorpos=args.cursorpos,
-                                   plugin=args.plugin,
-                                   alternate_language=args.alternate_language)
+        if args.extra_heartbeats:
+            heartbeats = json.loads(sys.stdin.read())
+            for heartbeat in heartbeats:
+                retval = process_heartbeat(args, configs, heartbeat)

-            project = args.project or args.alternate_project
-            branch = None
-            if args.entity_type == 'file':
-                project, branch = get_project_info(configs, args)
+        return retval

-            kwargs = vars(args)
-            kwargs['project'] = project
-            kwargs['branch'] = branch
-            kwargs['stats'] = stats
-            hostname = args.hostname or socket.gethostname()
-            kwargs['hostname'] = hostname
-            kwargs['timeout'] = args.timeout
-
-            status = send_heartbeat(**kwargs)
-            if status == SUCCESS:
-                return sync_offline_heartbeats(args, hostname)
-            else:
-                return status
-
-        else:
-            log.debug('File does not exist; ignoring this heartbeat.')
-            return SUCCESS
    except:
        log.traceback()
+        print(traceback.format_exc())
        return UNKNOWN_ERROR