Lots of fixes from code review.

3 files changed, 303 insertions, 181 deletions

diff --git a/python/subunit/_output.py b/python/subunit/_output.py
index bdea14f..49b5e81 100644
--- a/python/subunit/_output.py
+++ b/python/subunit/_output.py
@@ -12,25 +12,35 @@
 #  license you chose for the specific language governing permissions and
 #  limitations under that license.
 
+import datetime
+from functools import partial
 from optparse import (
     OptionGroup,
     OptionParser,
     OptionValueError,
 )
-import datetime
-from functools import partial
-from sys import stdin, stdout
-
-from testtools.compat import _b
+import sys
 
 from subunit.iso8601 import UTC
 from subunit.v2 import StreamResultToBytes
 
 
+_FINAL_ACTIONS = frozenset([
+    'exists',
+    'fail',
+    'skip',
+    'success',
+    'uxsuccess',
+    'xfail',
+])
+_ALL_ACTIONS = _FINAL_ACTIONS.union(['inprogress'])
+_CHUNK_SIZE=3670016 # 3.5 MiB
+
+
 def output_main():
     args = parse_arguments()
-    output = get_output_stream_writer()
-    generate_bytestream(args, output)
+    output = StreamResultToBytes(sys.stdout)
+    generate_stream_results(args, output)
     return 0
 
 
@@ -48,6 +58,7 @@ def parse_arguments(args=None, ParserClass=OptionParser):
         usage="subunit-output [-h] [status test_id] [options]",
     )
     parser.set_default('tags', None)
+    parser.set_default('test_id', None)
 
     status_commands = OptionGroup(
         parser,
@@ -55,21 +66,16 @@ def parse_arguments(args=None, ParserClass=OptionParser):
         "These options report the status of a test. TEST_ID must be a string "
             "that uniquely identifies the test."
     )
-    final_actions = 'exists fail skip success xfail uxsuccess'.split()
-    all_actions = final_actions + ['inprogress']
-    for action_name in all_actions:
-        final_text =  " This is a final state: No more status reports may "\
-            "be generated for this test id after this one."
-
+    for action_name in _ALL_ACTIONS:
         status_commands.add_option(
             "--%s" % action_name,
             nargs=1,
             action="callback",
-            callback=status_action,
+            callback=set_status_cb,
             callback_args=(action_name,),
             dest="action",
             metavar="TEST_ID",
-            help="Report a test status." + final_text if action_name in final_actions else ""
+            help="Report a test status."
         )
     parser.add_option_group(status_commands)
 
@@ -111,7 +117,7 @@ def parse_arguments(args=None, ParserClass=OptionParser):
         help="A comma-separated list of tags to associate with a test. This "
             "option may only be used with a status command.",
         action="callback",
-        callback=tags_action,
+        callback=set_tags_cb,
         default=[]
     )
 
@@ -124,7 +130,7 @@ def parse_arguments(args=None, ParserClass=OptionParser):
         if options.attach_file == '-':
             if not options.file_name:
                 options.file_name = 'stdin'
-            options.attach_file = stdin
+            options.attach_file = sys.stdin
         else:
             try:
                 options.attach_file = open(options.attach_file)
@@ -138,7 +144,7 @@ def parse_arguments(args=None, ParserClass=OptionParser):
     return options
 
 
-def status_action(option, opt_str, value, parser, status_name):
+def set_status_cb(option, opt_str, value, parser, status_name):
     if getattr(parser.values, "action", None) is not None:
         raise OptionValueError("argument %s: Only one status may be specified at once." % option)
 
@@ -148,60 +154,66 @@ def status_action(option, opt_str, value, parser, status_name):
     parser.values.test_id = parser.rargs.pop(0)
 
 
-def tags_action(option, opt_str, value, parser):
+def set_tags_cb(option, opt_str, value, parser):
     parser.values.tags = parser.rargs.pop(0).split(',')
 
 
-def get_output_stream_writer():
-    return StreamResultToBytes(stdout)
-
-
-def generate_bytestream(args, output_writer):
+def generate_stream_results(args, output_writer):
     output_writer.startTestRun()
+
     if args.attach_file:
-        write_chunked_file(
-            file_obj=args.attach_file,
-            test_id=args.test_id,
-            output_writer=output_writer,
-            mime_type=args.mimetype,
-        )
-    output_writer.status(
-        test_id=args.test_id,
-        test_status=args.action,
-        timestamp=create_timestamp(),
-        test_tags=args.tags,
-        )
-    output_writer.stopTestRun()
+        reader = partial(args.attach_file.read, _CHUNK_SIZE)
+        this_file_hunk = reader().encode('utf8')
+        next_file_hunk = reader().encode('utf8')
+
+    is_first_packet = True
+    is_last_packet = False
+    while not is_last_packet:
+
+        # XXX
+        def logme(*args, **kwargs):
+            print(args, kwargs)
+            output_writer.status(*args, **kwargs)
+        write_status = output_writer.status
+
+        if is_first_packet:
+            if args.attach_file:
+                # mimetype is specified on the first chunk only:
+                if args.mimetype:
+                    write_status = partial(write_status, mime_type=args.mimetype)
+            # tags are only written on the first packet:
+            if args.tags:
+                write_status = partial(write_status, test_tags=args.tags)
+            # timestamp is specified on the first chunk as well:
+            write_status = partial(write_status, timestamp=create_timestamp())
+            if args.action not in _FINAL_ACTIONS:
+                write_status = partial(write_status, test_status=args.action)
+            is_first_packet = False
+
+        if args.attach_file:
+            # filename might be overridden by the user
+            filename = args.file_name or args.attach_file.name
+            write_status = partial(write_status, file_name=filename, file_bytes=this_file_hunk)
+            if next_file_hunk == b'':
+                write_status = partial(write_status, eof=True)
+                is_last_packet = True
+            else:
+                this_file_hunk = next_file_hunk
+                next_file_hunk = reader().encode('utf8')
+        else:
+            is_last_packet = True
 
+        if args.test_id:
+            write_status = partial(write_status, test_id=args.test_id)
 
-def write_chunked_file(file_obj, output_writer, chunk_size=1024,
-                       mime_type=None, test_id=None, file_name=None):
-    reader = partial(file_obj.read, chunk_size)
+        if is_last_packet:
+            write_status = partial(write_status, eof=True)
+            if args.action in _FINAL_ACTIONS:
+                write_status = partial(write_status, test_status=args.action)
 
-    write_status = output_writer.status
-    if mime_type is not None:
-        write_status = partial(
-            write_status,
-            mime_type=mime_type
-        )
-    if test_id is not None:
-        write_status = partial(
-            write_status,
-            test_id=test_id
-        )
-    filename = file_name if file_name else file_obj.name
+        write_status()
 
-    for chunk in iter(reader, _b('')):
-        write_status(
-            file_name=filename,
-            file_bytes=chunk,
-            eof=False,
-        )
-    write_status(
-        file_name=filename,
-        file_bytes=_b(''),
-        eof=True,
-    )
+    output_writer.stopTestRun()
 
 
 def create_timestamp():
diff --git a/python/subunit/tests/test_output_filter.py b/python/subunit/tests/test_output_filter.py
index 21d8172..401ec08 100644
--- a/python/subunit/tests/test_output_filter.py
+++ b/python/subunit/tests/test_output_filter.py
@@ -13,6 +13,7 @@
 #  license you chose for the specific language governing permissions and
 #  limitations under that license.
 #
+
 import datetime
 from functools import partial
 from io import BytesIO, StringIO
@@ -20,9 +21,10 @@ import optparse
 import sys
 from tempfile import NamedTemporaryFile
 
+from contextlib import contextmanager
 from testscenarios import WithScenarios
 from testtools import TestCase
-from testtools.compat import _b, _u
+from testtools.compat import _u
 from testtools.matchers import (
     Equals,
     Matcher,
@@ -35,9 +37,10 @@ from testtools.testresult.doubles import StreamResult
 from subunit.iso8601 import UTC
 from subunit.v2 import StreamResultToBytes, ByteStreamToStreamResult
 from subunit._output import (
-    generate_bytestream,
+    _ALL_ACTIONS,
+    _FINAL_ACTIONS,
+    generate_stream_results,
     parse_arguments,
-    write_chunked_file,
 )
 import subunit._output as _o
 
@@ -67,15 +70,7 @@ class TestCaseWithPatchedStderr(TestCase):
 class TestStatusArgParserTests(WithScenarios, TestCaseWithPatchedStderr):
 
     scenarios = [
-        (cmd, dict(command=cmd, option='--' + cmd)) for cmd in (
-            'exists',
-            'fail',
-            'inprogress',
-            'skip',
-            'success',
-            'uxsuccess',
-            'xfail',
-        )
+        (cmd, dict(command=cmd, option='--' + cmd)) for cmd in _ALL_ACTIONS
     ]
 
     def test_can_parse_all_commands_with_test_id(self):
@@ -113,7 +108,7 @@ class TestStatusArgParserTests(WithScenarios, TestCaseWithPatchedStderr):
         self.assertThat(args.tags, Equals(["foo", "bar", "baz"]))
 
     def test_attach_file_with_hyphen_opens_stdin(self):
-        self.patch(_o, 'stdin', StringIO(_u("Hello")))
+        self.patch(_o.sys, 'stdin', StringIO(_u("Hello")))
         args = safe_parse_arguments(
             args=[self.option, "foo", "--attach-file", "-"]
         )
@@ -196,7 +191,7 @@ class ArgParserTests(TestCaseWithPatchedStderr):
 def get_result_for(commands):
     """Get a result object from *commands.
 
-    Runs the 'generate_bytestream' function from subunit._output after
+    Runs the 'generate_stream_results' function from subunit._output after
     parsing *commands as if they were specified on the command line. The
     resulting bytestream is then converted back into a result object and
     returned.
@@ -205,7 +200,7 @@ def get_result_for(commands):
 
     args = safe_parse_arguments(commands)
     output_writer = StreamResultToBytes(output_stream=stream)
-    generate_bytestream(args, output_writer)
+    generate_stream_results(args, output_writer)
 
     stream.seek(0)
 
@@ -215,143 +210,257 @@ def get_result_for(commands):
     return result
 
 
-class ByteStreamCompatibilityTests(WithScenarios, TestCase):
+@contextmanager
+def temp_file_contents(data):
+    """Create a temporary file on disk containing 'data'."""
+    with NamedTemporaryFile() as f:
+        f.write(data)
+        f.seek(0)
+        yield f
+
+
+class StatusStreamResultTests(WithScenarios, TestCase):
 
     scenarios = [
-        (s, dict(status=s, option='--' + s)) for s in (
-            'exists',
-            'fail',
-            'inprogress',
-            'skip',
-            'success',
-            'uxsuccess',
-            'xfail',
-        )
+        (s, dict(status=s, option='--' + s)) for s in _ALL_ACTIONS
     ]
 
     _dummy_timestamp = datetime.datetime(2013, 1, 1, 0, 0, 0, 0, UTC)
 
     def setUp(self):
-        super(ByteStreamCompatibilityTests, self).setUp()
+        super(StatusStreamResultTests, self).setUp()
         self.patch(_o, 'create_timestamp', lambda: self._dummy_timestamp)
+        self.test_id = self.getUniqueString()
+
+    def test_only_one_packet_is_generated(self):
+        result = get_result_for([self.option, self.test_id])
+        self.assertThat(
+            len(result._events),
+            Equals(1)
+        )
 
     def test_correct_status_is_generated(self):
-        result = get_result_for([self.option, 'foo'])
+        result = get_result_for([self.option, self.test_id])
 
         self.assertThat(
             result._events[0],
-            MatchesCall(
-                call='status',
-                test_id='foo',
-                test_status=self.status,
-                timestamp=self._dummy_timestamp,
-            )
+            MatchesStatusCall(test_status=self.status)
         )
 
-    def test_all_commands_accept_tags(self):
-        result = get_result_for([self.option, 'foo', '--tags', 'hello,world'])
+    def test_all_commands_generate_tags(self):
+        result = get_result_for([self.option, self.test_id, '--tags', 'hello,world'])
         self.assertThat(
             result._events[0],
-            MatchesCall(
-                call='status',
-                test_id='foo',
-                test_tags=set(['hello', 'world']),
-                timestamp=self._dummy_timestamp,
-            )
+            MatchesStatusCall(test_tags=set(['hello', 'world']))
         )
 
+    def test_all_commands_generate_timestamp(self):
+        result = get_result_for([self.option, self.test_id])
 
-class FileChunkingTests(WithScenarios, TestCase):
+        self.assertThat(
+            result._events[0],
+            MatchesStatusCall(timestamp=self._dummy_timestamp)
+        )
 
-    scenarios = [
-        ("With test_id", dict(test_id="foo")),
-        ("Without test_id", dict(test_id=None)),
-    ]
+    def test_all_commands_generate_correct_test_id(self):
+        result = get_result_for([self.option, self.test_id])
 
-    def _write_chunk_file(self, file_data, chunk_size=1024, mimetype=None, filename=None, test_id=None):
-        """Write file data to a subunit stream, get a StreamResult object."""
-        stream = BytesIO()
-        output_writer = StreamResultToBytes(output_stream=stream)
-
-        with NamedTemporaryFile() as f:
-            self._tmp_filename = f.name
-            f.write(file_data)
-            f.seek(0)
-
-            write_chunked_file(
-                file_obj=f,
-                output_writer=output_writer,
-                chunk_size=chunk_size,
-                mime_type=mimetype,
-                test_id=test_id,
-                file_name=filename,
+        self.assertThat(
+            result._events[0],
+            MatchesStatusCall(test_id=self.test_id)
+        )
+
+    def test_file_is_sent_in_single_packet(self):
+        with temp_file_contents(b"Hello") as f:
+            result = get_result_for([self.option, self.test_id, '--attach-file', f.name])
+
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(file_bytes=b'Hello', eof=True),
+                ])
             )
 
-        stream.seek(0)
+    def test_file_is_sent_with_test_id(self):
+        with temp_file_contents(b"Hello") as f:
+            result = get_result_for([self.option, self.test_id, '--attach-file', f.name])
 
-        case = ByteStreamToStreamResult(source=stream)
-        result = StreamResult()
-        case.run(result)
-        return result
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(test_id=self.test_id, file_bytes=b'Hello', eof=True),
+                ])
+            )
 
     def test_file_chunk_size_is_honored(self):
-        result = self._write_chunk_file(
-            file_data=_b("Hello"),
-            chunk_size=1,
-            test_id=self.test_id,
-        )
-        self.assertThat(
-            result._events,
-            MatchesListwise([
-                MatchesCall(call='status', test_id=self.test_id, file_bytes=_b('H'), mime_type=None, eof=False),
-                MatchesCall(call='status', test_id=self.test_id, file_bytes=_b('e'), mime_type=None, eof=False),
-                MatchesCall(call='status', test_id=self.test_id, file_bytes=_b('l'), mime_type=None, eof=False),
-                MatchesCall(call='status', test_id=self.test_id, file_bytes=_b('l'), mime_type=None, eof=False),
-                MatchesCall(call='status', test_id=self.test_id, file_bytes=_b('o'), mime_type=None, eof=False),
-                MatchesCall(call='status', test_id=self.test_id, file_bytes=_b(''), mime_type=None, eof=True),
+        with temp_file_contents(b"Hello") as f:
+            self.patch(_o, '_CHUNK_SIZE', 1)
+            result = get_result_for([self.option, self.test_id, '--attach-file', f.name])
+
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(test_id=self.test_id, file_bytes=b'H', eof=False),
+                    MatchesStatusCall(test_id=self.test_id, file_bytes=b'e', eof=False),
+                    MatchesStatusCall(test_id=self.test_id, file_bytes=b'l', eof=False),
+                    MatchesStatusCall(test_id=self.test_id, file_bytes=b'l', eof=False),
+                    MatchesStatusCall(test_id=self.test_id, file_bytes=b'o', eof=True),
+                ])
+            )
+
+    def test_file_mimetype_specified_once_only(self):
+        with temp_file_contents(b"Hi") as f:
+            self.patch(_o, '_CHUNK_SIZE', 1)
+            result = get_result_for([
+                self.option,
+                self.test_id,
+                '--attach-file',
+                f.name,
+                '--mimetype',
+                'text/plain',
             ])
-        )
 
-    def test_file_mimetype_is_honored(self):
-        result = self._write_chunk_file(
-            file_data=_b("SomeData"),
-            mimetype="text/plain",
-            test_id=self.test_id,
-        )
-        self.assertThat(
-            result._events,
-            MatchesListwise([
-                MatchesCall(call='status', test_id=self.test_id, file_bytes=_b('SomeData'), mime_type="text/plain"),
-                MatchesCall(call='status', test_id=self.test_id, file_bytes=_b(''), mime_type="text/plain"),
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(test_id=self.test_id, mime_type='text/plain', file_bytes=b'H', eof=False),
+                    MatchesStatusCall(test_id=self.test_id, mime_type=None, file_bytes=b'i', eof=True),
+                ])
+            )
+
+    def test_tags_specified_once_only(self):
+        with temp_file_contents(b"Hi") as f:
+            self.patch(_o, '_CHUNK_SIZE', 1)
+            result = get_result_for([
+                self.option,
+                self.test_id,
+                '--attach-file',
+                f.name,
+                '--tags',
+                'foo,bar',
             ])
-        )
 
-    def test_file_name_is_honored(self):
-        result = self._write_chunk_file(
-            file_data=_b("data"),
-            filename="/some/name",
-            test_id=self.test_id
-        )
-        self.assertThat(
-            result._events,
-            MatchesListwise([
-                MatchesCall(call='status', test_id=self.test_id, file_name='/some/name'),
-                MatchesCall(call='status', test_id=self.test_id, file_name='/some/name'),
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(test_id=self.test_id, test_tags=set(['foo', 'bar'])),
+                    MatchesStatusCall(test_id=self.test_id, test_tags=None),
+                ])
+            )
+
+    def test_timestamp_specified_once_only(self):
+        with temp_file_contents(b"Hi") as f:
+            self.patch(_o, '_CHUNK_SIZE', 1)
+            result = get_result_for([
+                self.option,
+                self.test_id,
+                '--attach-file',
+                f.name,
             ])
-        )
 
-    def test_default_filename_is_used(self):
-        result = self._write_chunk_file(file_data=_b("data"))
-        self.assertThat(
-            result._events,
-            MatchesListwise([
-                MatchesCall(call='status', file_name=self._tmp_filename),
-                MatchesCall(call='status', file_name=self._tmp_filename),
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(test_id=self.test_id, timestamp=self._dummy_timestamp),
+                    MatchesStatusCall(test_id=self.test_id, timestamp=None),
+                ])
+            )
+
+    def test_test_status_specified_once_only(self):
+        with temp_file_contents(b"Hi") as f:
+            self.patch(_o, '_CHUNK_SIZE', 1)
+            result = get_result_for([
+                self.option,
+                self.test_id,
+                '--attach-file',
+                f.name,
             ])
-        )
+
+            # 'inprogress' status should be on the first packet only, all other
+            # statuses should be on the last packet.
+            if self.status in _FINAL_ACTIONS:
+                first_call = MatchesStatusCall(test_id=self.test_id, test_status=None)
+                last_call = MatchesStatusCall(test_id=self.test_id, test_status=self.status)
+            else:
+                first_call = MatchesStatusCall(test_id=self.test_id, test_status=self.status)
+                last_call = MatchesStatusCall(test_id=self.test_id, test_status=None)
+            self.assertThat(
+                result._events,
+                MatchesListwise([first_call, last_call])
+            )
+
+    def test_filename_can_be_overridden(self):
+        with temp_file_contents(b"Hello") as f:
+            specified_file_name = self.getUniqueString()
+            result = get_result_for([
+                self.option,
+                self.test_id,
+                '--attach-file',
+                f.name,
+                '--file-name',
+                specified_file_name])
+
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(file_name=specified_file_name, file_bytes=b'Hello'),
+                ])
+            )
+
+    def test_file_name_is_used_by_default(self):
+        with temp_file_contents(b"Hello") as f:
+            result = get_result_for([self.option, self.test_id, '--attach-file', f.name])
+
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(file_name=f.name, file_bytes=b'Hello', eof=True),
+                ])
+            )
+
+
+class GlobalFileDataTests(TestCase):
+
+    def test_can_attach_file_without_test_id(self):
+        with temp_file_contents(b"Hello") as f:
+            result = get_result_for(['--attach-file', f.name])
+
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(test_id=None, file_bytes=b'Hello', eof=True),
+                ])
+            )
+
+    def test_file_name_is_used_by_default(self):
+        with temp_file_contents(b"Hello") as f:
+            result = get_result_for(['--attach-file', f.name])
+
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(file_name=f.name, file_bytes=b'Hello', eof=True),
+                ])
+            )
+
+    def test_filename_can_be_overridden(self):
+        with temp_file_contents(b"Hello") as f:
+            specified_file_name = self.getUniqueString()
+            result = get_result_for([
+                '--attach-file',
+                f.name,
+                '--file-name',
+                specified_file_name])
+
+            self.assertThat(
+                result._events,
+                MatchesListwise([
+                    MatchesStatusCall(file_name=specified_file_name, file_bytes=b'Hello'),
+                ])
+            )
 
 
-class MatchesCall(Matcher):
+class MatchesStatusCall(Matcher):
 
     _position_lookup = {
         'call': 0,
@@ -388,4 +497,4 @@ class MatchesCall(Matcher):
                 return Mismatch("Key %s is not present." % k)
 
     def __str__(self):
-        return "<MatchesCall %r>" % self._filters
+        return "<MatchesStatusCall %r>" % self._filters
diff --git a/setup.py b/setup.py
index 2f22300..1fe168d 100755
--- a/setup.py
+++ b/setup.py
@@ -9,9 +9,10 @@ except ImportError:
 else:
     extra = {
         'install_requires': [
+            'contextlib',
             'extras',
-            'testtools>=0.9.30',
             'testscenarios',
+            'testtools>=0.9.30',
         ]
     }