From 359e6b8b5f395d8af15706a3592a068a99ad7210 Mon Sep 17 00:00:00 2001
From: Clint Byrum <clint@fewbar.com>
Date: Wed, 11 Nov 2015 14:57:47 -0800
Subject: [PATCH] Add --output arg

This will allow us to print out the subuni stream on stdout, but write
the json itself on a seperate file. This way you can save the delta
result before processing it.
---
 os_performance_tools/collect.py            |  7 +++++-
 os_performance_tools/tests/test_collect.py | 28 ++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/os_performance_tools/collect.py b/os_performance_tools/collect.py
index 6da6ac7..dfae2f4 100644
--- a/os_performance_tools/collect.py
+++ b/os_performance_tools/collect.py
@@ -49,6 +49,8 @@ def main(argv=None, stdout=None):
                         help="Wrap the json output in a subunit stream. If an "
                         "argument is passed used that as the filename, "
                         "otherwise 'counters.json' will be used")
+    parser.add_argument('--output', help="Write JSON here. Does not disable "
+                        "stdout.")
     args = parser.parse_args(argv[1:])
     logging.basicConfig(
         format='%(asctime)-15s %(levelname)s %(threadName)s: %(message)s')
@@ -70,7 +72,6 @@ def main(argv=None, stdout=None):
     }
     if args.delta:
         collected = _delta.delta_with_file(args.delta, collected)
-
     content = json.dumps(collected, indent=1, sort_keys=True).encode('utf-8')
     if args.subunit is not None:
         file_name = args.subunit or 'counters.json'
@@ -82,6 +83,10 @@ def main(argv=None, stdout=None):
     else:
         stdout.write(content)
         stdout.write(b"\n")
+    if args.output:
+        with open(args.output, 'w') as output:
+            output.write(content)
+            output.write(b"\n")
 
 if __name__ == '__main__':
     main()
diff --git a/os_performance_tools/tests/test_collect.py b/os_performance_tools/tests/test_collect.py
index f7b4116..66c2f95 100644
--- a/os_performance_tools/tests/test_collect.py
+++ b/os_performance_tools/tests/test_collect.py
@@ -21,6 +21,7 @@ Tests for `os_performance_tools.collect`
 
 import json
 import mock
+import tempfile
 
 from os_performance_tools import collect
 from os_performance_tools.tests import base
@@ -85,3 +86,30 @@ class TestCollect(base.TestCase):
         self.assertTrue(isinstance(content, dict))
         self.assertIn('mysql', content)
         self.assertIn('queues', content)
+
+    @mock.patch('os_performance_tools.collectors.mysql.collect')
+    @mock.patch('os_performance_tools.collectors.queues.collect')
+    def test_collect_main_subunit_and_json(self, queues_mock, mysql_mock):
+        mysql_mock.return_value = {}
+        queues_mock.return_value = {}
+        with tempfile.NamedTemporaryFile() as tfile:
+            collect.main(
+                ['os-collect-counters', '--subunit', '--output', tfile.name],
+                self.stdout)
+            content = json.loads(tfile.read())
+            self.assertTrue(isinstance(content, dict))
+            self.assertIn('mysql', content)
+            self.assertIn('queues', content)
+        self.stdout.seek(0)
+        stream = subunit.ByteStreamToStreamResult(self.stdout)
+        result = StreamResult()
+        result.startTestRun()
+        try:
+            stream.run(result)
+        finally:
+            result.stopTestRun()
+        self.assertIsNotNone(result.counters_content)
+        content = json.loads(result.counters_content.decode('utf-8'))
+        self.assertTrue(isinstance(content, dict))
+        self.assertIn('mysql', content)
+        self.assertIn('queues', content)