From 45ede0d5adce61174e0b243c8342d3857de2a98d Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Wed, 20 Sep 2017 14:57:03 +0800
Subject: [PATCH] bugfix: once task_suite been executed, task in task_suite
 will be None, and can not call print_output function

---
 ate/cli.py  |  6 +++---
 ate/task.py | 20 ++++++++++++--------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/ate/cli.py b/ate/cli.py
index e8247fd8..0bf17b86 100644
--- a/ate/cli.py
+++ b/ate/cli.py
@@ -5,7 +5,7 @@ import sys
 from collections import OrderedDict
 
 from ate import __version__
-from ate.task import create_task
+from ate.task import TaskSuite
 
 import PyUnitReport
 
@@ -58,7 +58,7 @@ def main_ate():
     for testset_path in set(args.testset_paths):
 
         testset_path = testset_path.rstrip('/')
-        task_suite = create_task(testset_path)
+        task_suite = TaskSuite(testset_path)
 
         output_folder_name = os.path.basename(os.path.splitext(testset_path)[0])
         kwargs = {
@@ -78,7 +78,7 @@ def main_ate():
         if len(result.successes) != result.testsRun:
             subject = "FAILED"
 
-        for task in task_suite:
+        for task in task_suite.tasks:
             task.print_output()
 
     flag_code = 0 if subject == "SUCCESS" else 1
diff --git a/ate/task.py b/ate/task.py
index aa30fb86..a653dbff 100644
--- a/ate/task.py
+++ b/ate/task.py
@@ -42,16 +42,20 @@ class ApiTestSuite(unittest.TestSuite):
         output_variables_list = self.config_dict.get("output", [])
         self.test_runner.generate_output(output_variables_list)
 
-
-def create_task(testcase_path):
+class TaskSuite(unittest.TestSuite):
     """ create test task suite with specified testcase path.
         each task suite may include one or several test suite.
     """
-    task_suite = unittest.TestSuite()
-    testsets = testcase.load_testcases_by_path(testcase_path)
+    def __init__(self, testcase_path):
+        super(TaskSuite, self).__init__()
+        self.suite_list = []
+        testsets = testcase.load_testcases_by_path(testcase_path)
 
-    for testset in testsets:
-        suite = ApiTestSuite(testset)
-        task_suite.addTest(suite)
+        for testset in testsets:
+            suite = ApiTestSuite(testset)
+            self.addTest(suite)
+            self.suite_list.append(suite)
 
-    return task_suite
+    @property
+    def tasks(self):
+        return self.suite_list