[Findit] Refactoring WfTryJobData into BaseTryJobData, WfTryJobData, and FlakeTryJobData

appengine/findit/waterfall/test/identify_try_job_culprit_pipeline_test.py

Index: appengine/findit/waterfall/test/identify_try_job_culprit_pipeline_test.py
diff --git a/appengine/findit/waterfall/test/identify_try_job_culprit_pipeline_test.py b/appengine/findit/waterfall/test/identify_try_job_culprit_pipeline_test.py
index e3629d43ffb919d94c09f126b356879751ce4a31..5f6f6cbb3bdad50b85b5c5d340ee72ddba33ace5 100644
--- a/appengine/findit/waterfall/test/identify_try_job_culprit_pipeline_test.py
+++ b/appengine/findit/waterfall/test/identify_try_job_culprit_pipeline_test.py
@@ -434,6 +434,7 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     try_job.put()
     try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
     try_job_data.put()
 
     analysis = WfAnalysis.Create(master_name, builder_name, build_number)
@@ -468,6 +469,8 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
     }
 
     try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = WfTryJob.Create(
+        master_name, builder_name, build_number).key
     try_job_data.put()
 
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
@@ -550,8 +553,13 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
         'try_job_id': try_job_id,
     }
 
-    WfTryJobData.Create(try_job_id).put()
-    WfTryJob.Create(master_name, builder_name, build_number).put()
+    try_job = WfTryJob.Create(master_name, builder_name, build_number)
+    try_job.put()
+
+    try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
+    try_job_data.put()
+
     analysis = WfAnalysis.Create(master_name, builder_name, build_number)
     analysis.put()
 
@@ -576,10 +584,14 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
     build_number = 1
     try_job_id = '1'
 
-    WfTryJobData.Create(try_job_id).put()
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     try_job.status = analysis_status.RUNNING
     try_job.put()
+
+    try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
+    try_job_data.put()
+
     analysis = WfAnalysis.Create(master_name, builder_name, build_number)
     analysis.put()
 
@@ -608,11 +620,14 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
         'repo_name': 'chromium'
     }
 
-    WfTryJobData.Create(try_job_id).put()
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     try_job.status = analysis_status.RUNNING
     try_job.put()
 
+    try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
+    try_job_data.put()
+
     # Heuristic analysis already provided some results.
     analysis = WfAnalysis.Create(master_name, builder_name, build_number)
     analysis.result_status = result_status.FOUND_UNTRIAGED
@@ -655,10 +670,14 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
         'try_job_id': try_job_id
     }
 
-    WfTryJobData.Create(try_job_id).put()
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     try_job.status = analysis_status.RUNNING
     try_job.put()
+    
+    try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
+    try_job_data.put()
+
     analysis = WfAnalysis.Create(master_name, builder_name, build_number)
     analysis.put()
 
@@ -697,10 +716,14 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
         'try_job_id': try_job_id
     }
 
-    WfTryJobData.Create(try_job_id).put()
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     try_job.status = analysis_status.RUNNING
     try_job.put()
+
+    try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
+    try_job_data.put()
+
     analysis = WfAnalysis.Create(master_name, builder_name, build_number)
     analysis.put()
 
@@ -783,11 +806,15 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
         'try_job_id': try_job_id
     }
 
-    WfTryJobData.Create(try_job_id).put()
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     try_job.status = analysis_status.RUNNING
     try_job.test_results = [test_result]
     try_job.put()
+
+    try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
+    try_job_data.put()
+
     analysis = WfAnalysis.Create(master_name, builder_name, build_number)
     analysis.put()
 
@@ -947,9 +974,6 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
         },
     }
 
-    try_job_data = WfTryJobData.Create(try_job_id)
-    try_job_data.put()
-
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     try_job.status = analysis_status.RUNNING
     try_job.compile_results = [{
@@ -960,9 +984,12 @@ class IdentifyTryJobCulpritPipelineTest(testing.AppengineTestCase):
         },
         'try_job_id': try_job_id,
     }]
-
     try_job.put()
 
+    try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
+    try_job_data.put()
+
     pipeline = IdentifyTryJobCulpritPipeline()
     pipeline.run(master_name, builder_name, build_number, [revision],
                  failure_type.COMPILE, '1', compile_result)