Clean up CFG test suite

PiperOrigin-RevId: 197059149

1 files changed, 50 insertions, 37 deletions

diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/cfg_test.py b/tensorflow/contrib/autograph/pyct/static_analysis/cfg_test.py
index 8d723ce09d..fc07fa3447 100644
--- a/tensorflow/contrib/autograph/pyct/static_analysis/cfg_test.py
+++ b/tensorflow/contrib/autograph/pyct/static_analysis/cfg_test.py
@@ -115,20 +115,27 @@ class CFGTest(test.TestCase):
     if_body = body[0].body
     self._check_anno_matches(if_body[0], 'defined_out', ('x', 'y'))
 
-  # TODO(alexbw): b/73926938 split this test up
-  def test_live(self):
+  def _get_live_annotated_fnbody(self, f):
+    node, ctx = self._parse_and_analyze(f, {})
+    cfg.run_analyses(node, cfg.Liveness(ctx))
+    body = node.body[0].body
+    return body
 
-    def get_live_annotated_fnbody(f):
-      node, ctx = self._parse_and_analyze(f, {})
-      cfg.run_analyses(node, cfg.Liveness(ctx))
-      body = node.body[0].body
-      return body
+  def test_live_straightline(self):
 
     def f1(x):
       a = g(x)  # pylint: disable=undefined-variable
       b = h(a)  # pylint: disable=undefined-variable, unused-variable
       return x
 
+    body = self._get_live_annotated_fnbody(f1)
+    self._check_anno_matches(body[1], 'live_in', ('a', 'h', 'x'))
+    self._check_anno_matches(body[2], 'live_in', ('x'))
+    self._check_anno_matches(body[0], 'live_in', ('g', 'h', 'x'))
+    self._check_anno_matches(body[2], 'live_out', ())
+
+  def test_live_stacked_conds_with_else(self):
+
     def f2(x, a):  # pylint: disable=unused-argument
       if a > 0:  # x should not be live
         x = 0
@@ -137,6 +144,12 @@ class CFGTest(test.TestCase):
       else:
         x = 2
 
+    body = self._get_live_annotated_fnbody(f2)
+    self._check_anno_matches(body[0], 'live_in', ('a'))
+    self._check_anno_matches(body[1], 'live_in', ('a'))
+
+  def test_live_stacked_conds(self):
+
     def f3(x, a):
       if a > 0:  # x and a should be live
         x = 0
@@ -144,58 +157,58 @@ class CFGTest(test.TestCase):
         x = 1
       return x  # x should be live
 
+    body = self._get_live_annotated_fnbody(f3)
+    self._check_anno_matches(body[0], 'live_in', ('a', 'x'))
+    self._check_anno_matches(body[1], 'live_in', ('a', 'x'))
+    self._check_anno_matches(body[2], 'live_in', ('x'))
+
+  def test_live_possibly_unused_cond(self):
+
     def f4(x, a):
       if a > 0:  # x should be live
         x = 0
       x += 1
 
+    body = self._get_live_annotated_fnbody(f4)
+    self._check_anno_matches(body[0], 'live_in', ('x', 'a'))
+    self._check_anno_matches(body[1], 'live_in', ('x'))
+
+  def test_live_attribute_in_cond(self):
+
     def f5(x, a):
       if a > 0:  # x.y should be live
         x.y = 0
       return x.y
 
+    body = self._get_live_annotated_fnbody(f5)
+    self._check_anno_matches(body[0], 'live_in', ('x', 'x.y', 'a'))
+
+  def test_live_noop(self):
+
     def f6(x):
       return x  # should this cause x.* to be live?
 
+    body = self._get_live_annotated_fnbody(f6)
+    self._check_anno_matches(body[0], 'live_in', ('x'))
+
+  def test_live_loop(self):
+
     def f7(x, n):
       for i in range(n):
         x += i
       return x
 
-    def f8(x, f):
-      with f:
-        x += 1
-
-    body = get_live_annotated_fnbody(f1)
-    self._check_anno_matches(body[1], 'live_in', ('a', 'h', 'x'))
-    self._check_anno_matches(body[2], 'live_in', ('x'))
-    self._check_anno_matches(body[0], 'live_in', ('g', 'h', 'x'))
-    self._check_anno_matches(body[2], 'live_out', ())
-
-    body = get_live_annotated_fnbody(f2)
-    self._check_anno_matches(body[0], 'live_in', ('a'))
-    self._check_anno_matches(body[1], 'live_in', ('a'))
-
-    body = get_live_annotated_fnbody(f3)
-    self._check_anno_matches(body[0], 'live_in', ('a', 'x'))
-    self._check_anno_matches(body[1], 'live_in', ('a', 'x'))
-    self._check_anno_matches(body[2], 'live_in', ('x'))
-
-    body = get_live_annotated_fnbody(f4)
-    self._check_anno_matches(body[0], 'live_in', ('x', 'a'))
+    body = self._get_live_annotated_fnbody(f7)
+    self._check_anno_matches(body[0], 'live_in', ('x', 'n', 'range'))
     self._check_anno_matches(body[1], 'live_in', ('x'))
 
-    body = get_live_annotated_fnbody(f5)
-    self._check_anno_matches(body[0], 'live_in', ('x', 'x.y', 'a'))
+  def test_live_context_manager(self):
 
-    body = get_live_annotated_fnbody(f6)
-    self._check_anno_matches(body[0], 'live_in', ('x'))
-
-    body = get_live_annotated_fnbody(f7)
-    self._check_anno_matches(body[0], 'live_in', ('x', 'n', 'range'))
-    self._check_anno_matches(body[1], 'live_in', ('x'))
+    def f8(x, f):
+      with f:
+        x += 1
 
-    body = get_live_annotated_fnbody(f8)
+    body = self._get_live_annotated_fnbody(f8)
     self._check_anno_matches(body[0], 'live_in', ('f', 'x'))
 
   def test_node_equality(self):