Promote locals in tests to module scope (google#90)

Local-scoped classes won't be pickle-able following an internal change.

Author: Northbadge

compiler_opt/distributed/local/local_worker_manager_test.py

@@ -24,30 +24,46 @@
 from tf_agents.system import system_multiprocessing as multiprocessing
 
 
-class LocalWorkerManagerTest(absltest.TestCase):
+class JobNormal(Worker):
+  """Test worker."""
 
-  def test_pool(self):
+  def __init__(self):
+    self._token = 0
+
+  @classmethod
+  def is_priority_method(cls, method_name: str) -> bool:
+    return method_name == 'priority_method'
+
+  def priority_method(self):
+    return f'priority {self._token}'
+
+  def get_token(self):
+    return self._token
 
-    class Job(Worker):
-      """Test worker."""
+  def set_token(self, value):
+    self._token = value
 
-      def __init__(self):
-        self._token = 0
 
-      @classmethod
-      def is_priority_method(cls, method_name: str) -> bool:
-        return method_name == 'priority_method'
+class JobFail(Worker):
 
-      def priority_method(self):
-        return f'priority {self._token}'
+  def __init__(self, wont_be_passed):
+    self._arg = wont_be_passed
 
-      def get_token(self):
-        return self._token
+  def method(self):
+    return self._arg
 
-      def set_token(self, value):
-        self._token = value
 
-    with local_worker_manager.LocalWorkerPool(Job, 2) as pool:
+class JobSlow(Worker):
+
+  def method(self):
+    time.sleep(3600)
+
+
+class LocalWorkerManagerTest(absltest.TestCase):
+
+  def test_pool(self):
+
+    with local_worker_manager.LocalWorkerPool(JobNormal, 2) as pool:
       p1 = pool[0]
       p2 = pool[1]
       set_futures = [p1.set_token(1), p2.set_token(2)]
@@ -66,28 +82,15 @@ def set_token(self, value):
 
   def test_failure(self):
 
-    class Job(Worker):
-
-      def __init__(self, wont_be_passed):
-        self._arg = wont_be_passed
-
-      def method(self):
-        return self._arg
-
-    with local_worker_manager.LocalWorkerPool(Job, 2) as pool:
+    with local_worker_manager.LocalWorkerPool(JobFail, 2) as pool:
       with self.assertRaises(concurrent.futures.CancelledError):
         # this will fail because we didn't pass the arg to the ctor, so the
         # worker hosting process will crash.
         pool[0].method().result()
 
   def test_worker_crash_while_waiting(self):
 
-    class Job(Worker):
-
-      def method(self):
-        time.sleep(3600)
-
-    with local_worker_manager.LocalWorkerPool(Job, 2) as pool:
+    with local_worker_manager.LocalWorkerPool(JobSlow, 2) as pool:
       p = pool[0]
       f = p.method()
       self.assertFalse(f.done())

compiler_opt/rl/local_data_collector_test.py

@@ -89,14 +89,15 @@ def collect_data(self, module_spec, tf_policy_path, reward_stat):
         sequence_examples=[], reward_stats={}, rewards=[], keys=[])
 
 
-class LocalDataCollectorTest(tf.test.TestCase):
+class MyRunner(compilation_runner.CompilationRunner):
 
-  def test_local_data_collector(self):
+  def collect_data(self, *args, **kwargs):
+    return mock_collect_data(*args, **kwargs)
 
-    class MyRunner(compilation_runner.CompilationRunner):
 
-      def collect_data(self, *args, **kwargs):
-        return mock_collect_data(*args, **kwargs)
+class LocalDataCollectorTest(tf.test.TestCase):
+
+  def test_local_data_collector(self):
 
     def create_test_iterator_fn():