Add support for attaching files.

2 files changed, 70 insertions, 5 deletions

diff --git a/python/subunit/_output.py b/python/subunit/_output.py
index b3a5bba..43097e6 100644
--- a/python/subunit/_output.py
+++ b/python/subunit/_output.py
@@ -15,6 +15,7 @@
 
 from argparse import ArgumentParser
 import datetime
+from functools import partial
 from sys import stdout
 
 from subunit.v2 import StreamResultToBytes
@@ -46,6 +47,11 @@ def parse_arguments(args=None, ParserClass=ArgumentParser):
         "test_id",
         help="A string that uniquely identifies this test."
     )
+    common_args.add_argument(
+        "--attach-file",
+        type=file,
+        help="Attach a file to the result stream for this test."
+    )
     sub_parsers = parser.add_subparsers(dest="action")
 
     final_state = "This is a final action: No more actions may be generated " \
@@ -101,6 +107,8 @@ def get_output_stream_writer():
 
 def generate_bytestream(args, output_writer):
     output_writer.startTestRun()
+    if args.attach_file:
+        write_chunked_file(args.attach_file, args.test_id, output_writer)
     output_writer.status(
         test_id=args.test_id,
         test_status=translate_command_name(args.action),
@@ -109,6 +117,23 @@ def generate_bytestream(args, output_writer):
     output_writer.stopTestRun()
 
 
+def write_chunked_file(file_obj, test_id, output_writer, chunk_size=1024):
+    reader = partial(file_obj.read, chunk_size)
+    for chunk in iter(reader, ''):
+        output_writer.status(
+            test_id=test_id,
+            file_name=file_obj.name,
+            file_bytes=chunk,
+            eof=False,
+        )
+    output_writer.status(
+            test_id=test_id,
+            file_name=file_obj.name,
+            file_bytes='',
+            eof=True,
+        )
+
+
 _ZERO = datetime.timedelta(0)
 
 
diff --git a/python/subunit/tests/test_output_filter.py b/python/subunit/tests/test_output_filter.py
index c9059df..9d530c5 100644
--- a/python/subunit/tests/test_output_filter.py
+++ b/python/subunit/tests/test_output_filter.py
@@ -20,10 +20,13 @@ from collections import namedtuple
 import datetime
 from functools import partial
 from io import BytesIO
+from tempfile import NamedTemporaryFile
 from testtools import TestCase
 from testtools.matchers import (
     Equals,
+    IsInstance,
     Matcher,
+    MatchesListwise,
     Mismatch,
 )
 from testtools.testresult.doubles import StreamResult
@@ -34,6 +37,7 @@ from subunit._output import (
     parse_arguments,
     translate_command_name,
     utc,
+    write_chunked_file,
 )
 import subunit._output as _o
 
@@ -71,11 +75,13 @@ class OutputFilterArgumentTests(TestCase):
             self.assertThat(translate_command_name(command), Equals(command))
 
     def test_all_commands_parse_file_attachment(self):
-        for command in self._all_supported_commands:
-            args = safe_parse_arguments(
-                args=[command, 'foo', '--attach-file', '/some/path']
-            )
-            self.assertThat(args.attach_file, Equals('/some/path'))
+        with NamedTemporaryFile() as tmp_file:
+            for command in self._all_supported_commands:
+                args = safe_parse_arguments(
+                    args=[command, 'foo', '--attach-file', tmp_file.name]
+                )
+                self.assertThat(args.attach_file, IsInstance(file))
+                self.assertThat(args.attach_file.name, Equals(tmp_file.name))
 
 
 class ByteStreamCompatibilityTests(TestCase):
@@ -185,6 +191,40 @@ class ByteStreamCompatibilityTests(TestCase):
         )
 
 
+class FileChunkingTests(TestCase):
+
+    def _write_chunk_file(self, file_data, chunk_size):
+        """Write chunked data to a subunit stream, return a StreamResult object."""
+        stream = BytesIO()
+        output_writer = StreamResultToBytes(output_stream=stream)
+
+        with NamedTemporaryFile() as f:
+            f.write(file_data)
+            f.seek(0)
+
+            write_chunked_file(f, 'foo_test', output_writer, chunk_size)
+
+        stream.seek(0)
+
+        case = ByteStreamToStreamResult(source=stream)
+        result = StreamResult()
+        case.run(result)
+        return result
+
+    def test_file_chunk_size_is_honored(self):
+        result = self._write_chunk_file("Hello", 1)
+        self.assertThat(
+            result._events,
+            MatchesListwise([
+                MatchesCall(call='status', file_bytes='H', eof=False),
+                MatchesCall(call='status', file_bytes='e', eof=False),
+                MatchesCall(call='status', file_bytes='l', eof=False),
+                MatchesCall(call='status', file_bytes='l', eof=False),
+                MatchesCall(call='status', file_bytes='o', eof=False),
+                MatchesCall(call='status', file_bytes='', eof=True),
+            ])
+        )
+
 class MatchesCall(Matcher):
 
     _position_lookup = {