From 8206b38d3a256017aac0049f65d5b46a8e02375c Mon Sep 17 00:00:00 2001
From: Matthew Treinish <mtreinish@kortar.org>
Date: Wed, 27 May 2015 16:18:58 -0400
Subject: [PATCH] Add support to subunit2sql for storing attachments

This commit adds support to the subunit2sql script for storing
attachments from the subunit stream into the attachments table in
the database.

Change-Id: I5ccc0872a5a4bfca1eecd25059850c4a5c4f6b2b
---
 subunit2sql/db/api.py       | 31 +++++++++++++++++++++++++++++++
 subunit2sql/read_subunit.py | 16 ++++++++++++++--
 subunit2sql/shell.py        | 13 ++++++++++---
 3 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/subunit2sql/db/api.py b/subunit2sql/db/api.py
index a793fec..d16be24 100644
--- a/subunit2sql/db/api.py
+++ b/subunit2sql/db/api.py
@@ -746,3 +746,34 @@ def get_test_counts_in_date_range(test_id, start_date=None, stop_date=None,
     count_dict['failure'] = fail_query.count()
     count_dict['skips'] = skip_query.count()
     return count_dict
+
+
+def add_test_run_attachments(attach_dict, test_run_id, session=None):
+    """Add attachments a specific test run.
+
+    This method will take a dictionary and store key blob pair attachments in
+    the DB associated with the specified test_run.
+
+    :param dict attachments_dict: a dictionary which will generate a separate
+                                  key blob pair row associated with the
+                                  test_run_id
+    :param str test_run_id: the uuid of the test_run to update. (value of the
+                            id column for the row to be updated)
+    :param session: optional session object if one isn't provided a new session
+                    will be acquired for the duration of this operation
+
+    :return list: The list of created attachment objects
+    :rtype: subunit2sql.models.Attachments
+    """
+
+    session = session or get_session()
+    attachments = []
+    for label, attach in attach_dict.items():
+        attachment = models.Attachments()
+        attachment.label = label
+        attachment.attachment = attach
+        attachment.test_run_id = test_run_id
+        with session.begin():
+            session.add(attachment)
+        attachments.append(attachment)
+    return attachments
diff --git a/subunit2sql/read_subunit.py b/subunit2sql/read_subunit.py
index 8e761eb..0d6bd8f 100644
--- a/subunit2sql/read_subunit.py
+++ b/subunit2sql/read_subunit.py
@@ -32,7 +32,7 @@ def get_duration(start, end):
 
 class ReadSubunit(object):
 
-    def __init__(self, stream_file):
+    def __init__(self, stream_file, attachments=False):
         self.stream_file = stream_file
         self.stream = subunit.ByteStreamToStreamResult(stream_file)
         starts = testtools.StreamResult()
@@ -41,6 +41,7 @@ class ReadSubunit(object):
             self.parse_outcome))
         self.result = testtools.CopyStreamResult([starts, outcomes, summary])
         self.results = {}
+        self.attachments = attachments
 
     def get_results(self):
         self.result.startTestRun()
@@ -51,6 +52,13 @@ class ReadSubunit(object):
         self.results['run_time'] = self.run_time()
         return self.results
 
+    def get_attachments(self, test):
+        attach_dict = {}
+        for name, detail in test['details'].items():
+            name = name.split(':')[0]
+            attach_dict[name] = detail
+        return attach_dict
+
     def parse_outcome(self, test):
         metadata = {}
         status = test['status']
@@ -66,11 +74,15 @@ class ReadSubunit(object):
         if name == 'process-returncode':
             return
         timestamps = test['timestamps']
+        attachment_dict = {}
+        if self.attachments:
+            attachment_dict = self.get_attachments(test)
         self.results[name] = {
             'status': status,
             'start_time': timestamps[0],
             'end_time': timestamps[1],
-            'metadata': metadata
+            'metadata': metadata,
+            'attachments': attachment_dict
         }
         self.stream_file.flush()
 
diff --git a/subunit2sql/shell.py b/subunit2sql/shell.py
index 8a7510e..7d889cb 100644
--- a/subunit2sql/shell.py
+++ b/subunit2sql/shell.py
@@ -33,6 +33,8 @@ SHELL_OPTS = [
                 help='Dict of metadata about the run(s)'),
     cfg.StrOpt('artifacts', short='a', default=None,
                help='Location of run artifacts'),
+    cfg.BoolOpt('store_attachments', short='s', default=False,
+                help='Store attachments from subunit streams in the DB'),
     cfg.StrOpt('run_id', short='i', default=None,
                help='Run id to use for the specified subunit stream, can only'
                     ' be used if a single stream is provided')
@@ -142,6 +144,9 @@ def process_results(results):
         if results[test]['metadata']:
             api.add_test_run_metadata(results[test]['metadata'], test_run.id,
                                       session)
+        if results[test]['attachments']:
+            api.add_test_run_attachments(results[test]['attachment'],
+                                         test_run.id, session)
     session.close()
 
 
@@ -152,10 +157,12 @@ def main():
         if len(CONF.subunit_files) > 1 and CONF.run_id:
             print("You can not specify a run id for adding more than 1 stream")
             return 3
-        streams = [subunit.ReadSubunit(open(s, 'r')) for s in
-                   CONF.subunit_files]
+        streams = [subunit.ReadSubunit(open(s, 'r'),
+                                       attachments=CONF.store_attachments)
+                   for s in CONF.subunit_files]
     else:
-        streams = [subunit.ReadSubunit(sys.stdin)]
+        streams = [subunit.ReadSubunit(sys.stdin,
+                                       attachments=CONF.store_attachments)]
     for stream in streams:
         process_results(stream.get_results())