text: """Tests for dfs_chase_gen."""
  # Access to protected class members is common for unit tests.
  # pylint: disable=protected-access
  
  from frontend.proto import cfg_pb2
  from frontend.cfg_generator import dfs_chase_gen
  import unittest
  
  
  class DirectCallDFSChaseGenTest(unittest.TestCase):
  
      def setUp(self):
          self.depth = 3
          self.branch_probability = 0.5
          self.gen = dfs_chase_gen.DFSChaseGenerator(self.depth, False,
                                                     self.branch_probability)
  
      def test_generate_function_tree(self):
          self.gen._generate_function_tree()
          # A full binary tree of depth N has 2^n-1 nodes. function_tree only
          # includes functions that call others, so the leaves in the last layer
          # of the tree are excluded. Hence the expected number of nodes is
          # 2^(n-1) - 1.
          self.assertEqual(len(self.gen._function_tree), 2**(self.depth - 1) - 1)
          for children in self.gen._function_tree.values():
              if children:
                  self.assertEqual(len(children), 2)
  
      def test_generate_conditional_branch_code_blocks(self):
          callees = [2, 3]  # Function IDs.
          codeblock, ft_block, ft_block_ret, taken_block, taken_block_ret = \
              self.gen._generate_conditional_branch_code_blocks(
                  callees, self.branch_probability)
          self.assertEqual(codeblock.terminator_branch.type,
                           cfg_pb2.Branch.BranchType.CONDITIONAL_DIRECT)
          self.assertEqual(codeblock.terminator_branch.targets, [taken_block.id])
          self.assertEqual(codeblock.terminator_branch.taken_probability,
                           [self.branch_probability])
  
          self.assertEqual(taken_block.terminator_branch.type,
                           cfg_pb2.Branch.BranchType.DIRECT_CALL)
          self.assertEqual(taken_block.terminator_branch.targets, [callees[0]])
  
          self.assertEqual(ft_block.terminator_branch.type,
                           cfg_pb2.Branch.BranchType.DIRECT_CALL)
          self.assertEqual(ft_block.terminator_branch.targets, [callees[1]])
  
          self.assertEqual(taken_block_ret.terminator_branch.type,
                           cfg_pb2.Branch.BranchType.RETURN)
          self.assertEqual(ft_block_ret.terminator_branch.type,
                           cfg_pb2.Branch.BranchType.RETURN)
  
      def test_generate_functions(self):
          self.gen._generate_function_tree()
          self.gen._generate_functions()
          self.assertEqual(len(self.gen._functions), 2**(self.depth) - 1)
          for func_id, func in self.gen._functions.items():
              if func_id in self.gen._function_tree:
                  # Functions that call other functions have 5 call blocks.
                  self.assertEqual(len(func.instructions), 5)
              else:
                  # Leaf functions have just one.
                  self.assertEqual(len(func.instructions), 1)
  
      def test_generate_cfg(self):
          cfg = self.gen.generate_cfg()
          self.assertEqual(cfg.entry_point_function, self.gen._root_func)
          self.assertEqual(len(cfg.functions), len(self.gen._functions))
          for got_func in cfg.functions:
              self.assertEqual(got_func, self.gen._functions[got_func.id])
          # One CodeBlockBody per function signature and the shared code block
          # body for all functions.
          expected_codeblock_bodies = len(self.gen._functions) + 1
          self.assertEqual(len(cfg.code_block_bodies), expected_codeblock_bodies)
  
  
  class IndirectCallDFSChaseGenTest(unittest.TestCase):
  
      def setUp(self):
          self.depth = 3
          self.branch_probability = 0.5
          self.gen = dfs_chase_gen.DFSChaseGenerator(self.depth, True,
                                                     self.branch_probability)
  
      def test_generate_indirect_call_code_block(self):
          callee_funcs = [2, 3]
          callee_probability = 0.6
          block = self.gen._generate_indirect_call_code_block(
              callee_funcs, callee_probability)
          self.assertEqual(block.terminator_branch.type,
                           cfg_pb2.Branch.BranchType.INDIRECT_CALL)
          self.assertEqual(block.terminator_branch.targets, callee_funcs)
          self.assertAlmostEqual(block.terminator_branch.taken_probability[0],
                                 0.6)
          self.assertAlmostEqual(block.terminator_branch.taken_probability[1],
                                 0.4)
  
  
  if __name__ == '__main__':
      unittest.main()
  