src/third_party/py_trace_event/src/trace_event_impl/multiprocessing_shim_test.py

   1 # Copyright 2011 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4 import multiprocessing
   5 import tempfile
   6 import time
   7 import unittest
   8 from .log import *
   9 from .trace_test import *
  10
  11 import os
  12
  13 def DoWork():
  14   """
  15   Sadly, you can't @trace toplevel functions, as it prevents them
  16   from being called in the child process. :(
  17
  18   So, we wrap the function of interest.
  19   """
  20   def do_work():
  21     trace_begin("do_work")
  22     time.sleep(0.25)
  23     trace_end("do_work")
  24   do_work()
  25
  26 def AssertTracingEnabled():
  27   assert trace_is_enabled()
  28
  29 def AssertTracingDisabled():
  30   assert not trace_is_enabled()
  31
  32 def TryToDisableTracing():
  33   trace_disable();
  34
  35 class MultiprocessingShimTest(TraceTest):
  36   def test_shimmed(self):
  37     p = multiprocessing.Process()
  38     self.assertTrue(hasattr(p, "_shimmed_by_trace_event"))
  39
  40   def test_trace_enable_throws_in_child(self):
  41     def work():
  42       trace_begin("work")
  43       p = multiprocessing.Pool(1)
  44       self.assertRaises(Exception, lambda: p.apply(TryToDisableTracing, ()))
  45       p.close()
  46       p.terminate()
  47       p.join()
  48       trace_end("work")
  49     res = self.go(work)
  50
  51   def test_trace_enabled_in_child(self):
  52     def work():
  53       trace_begin("work")
  54       p = multiprocessing.Pool(1)
  55       p.apply(AssertTracingEnabled, ())
  56       p.close()
  57       p.terminate()
  58       p.join()
  59       trace_end("work")
  60     res = self.go(work)
  61
  62   def test_one_func(self):
  63     def work():
  64       trace_begin("work")
  65       p = multiprocessing.Pool(1)
  66       p.apply(DoWork, ())
  67       p.close()
  68       p.terminate()
  69       p.join()
  70       trace_end("work")
  71     res = self.go(work)
  72     work_events = res.findByName('work')
  73     do_work_events = res.findByName('do_work')
  74     self.assertEquals(2, len(work_events))
  75     self.assertEquals(2, len(do_work_events))
  76