Refactor test ESWorker to separate class

compiler_opt/es/blackbox_evaluator_test.py

@@ -20,7 +20,7 @@
 
 from compiler_opt.distributed.local import local_worker_manager
 from compiler_opt.rl import corpus
-from compiler_opt.es import blackbox_learner_test
+from compiler_opt.es import blackbox_test_utils
 from compiler_opt.es import blackbox_evaluator
 
 
@@ -29,7 +29,7 @@ class BlackboxEvaluatorTests(absltest.TestCase):
 
   def test_sampling_get_results(self):
     with local_worker_manager.LocalWorkerPoolManager(
-        blackbox_learner_test.ESWorker, count=3, arg='', kwarg='') as pool:
+        blackbox_test_utils.ESWorker, count=3, arg='', kwarg='') as pool:
       perturbations = [b'00', b'01', b'10']
       evaluator = blackbox_evaluator.SamplingBlackboxEvaluator(None, 5, 5, None)
       # pylint: disable=protected-access

compiler_opt/es/blackbox_learner_test.py

@@ -18,40 +18,19 @@
 from absl.testing import absltest
 import gin
 import tempfile
-from typing import List
 import numpy as np
 import numpy.typing as npt
 import tensorflow as tf
 from tf_agents.networks import actor_distribution_network
 from tf_agents.policies import actor_policy
 
-from compiler_opt.distributed import worker
 from compiler_opt.distributed.local import local_worker_manager
 from compiler_opt.es import blackbox_learner, policy_utils
 from compiler_opt.es import blackbox_optimizers
 from compiler_opt.rl import corpus, inlining, policy_saver, registry
 from compiler_opt.rl.inlining import config as inlining_config
 from compiler_opt.es import blackbox_evaluator
-
-
-@gin.configurable
-class ESWorker(worker.Worker):
-  """Temporary placeholder worker.
-  Each time a worker is called, the function value
-  it will return increases."""
-
-  def __init__(self, arg, *, kwarg):
-    self._arg = arg
-    self._kwarg = kwarg
-    self.function_value = 0.0
-
-  def compile(self, policy: policy_saver.Policy,
-              samples: List[corpus.ModuleSpec]) -> float:
-    if policy and samples:
-      self.function_value += 1.0
-      return self.function_value
-    else:
-      return 0.0
+from compiler_opt.es import blackbox_test_utils
 
 
 class BlackboxLearnerTests(absltest.TestCase):
@@ -161,7 +140,7 @@ def test_prune_skipped_perturbations(self):
 
   def test_run_step(self):
     with local_worker_manager.LocalWorkerPoolManager(
-        ESWorker, count=3, arg='', kwarg='') as pool:
+        blackbox_test_utils.ESWorker, count=3, arg='', kwarg='') as pool:
       self._learner.run_step(pool)  # pylint: disable=protected-access
       # expected length calculated from expected shapes of variables
       self.assertEqual(len(self._learner.get_model_weights()), 17154)

compiler_opt/es/blackbox_test_utils.py

@@ -0,0 +1,43 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Test facilities for Blackbox classes."""
+
+from typing import List
+
+import gin
+
+from compiler_opt.distributed import worker
+from compiler_opt.rl import corpus
+from compiler_opt.rl import policy_saver
+
+
+@gin.configurable
+class ESWorker(worker.Worker):
+  """Temporary placeholder worker.
+  Each time a worker is called, the function value
+  it will return increases."""
+
+  def __init__(self, arg, *, kwarg):
+    self._arg = arg
+    self._kwarg = kwarg
+    self.function_value = 0.0
+
+  def compile(self, policy: policy_saver.Policy,
+              samples: List[corpus.ModuleSpec]) -> float:
+    if policy and samples:
+      self.function_value += 1.0
+      return self.function_value
+    else:
+      return 0.0