Handle 3.8->3.9 unittest differences (#38)

Author: mtrofin

compiler_opt/rl/data_collector_test.py

@@ -15,7 +15,7 @@
 """Tests for data_collector."""
 
 # pylint: disable=protected-access
-
+import sys
 from unittest import mock
 
 from absl.testing import absltest
@@ -28,7 +28,12 @@ class DataCollectorTest(absltest.TestCase):
   def test_build_distribution_monitor(self):
     data = [3, 2, 1]
     monitor_dict = data_collector.build_distribution_monitor(data)
-    self.assertEqual(monitor_dict, monitor_dict | {'mean': 2, 'p_0.1': 1})
+    reference_dict = {'mean': 2, 'p_0.1': 1}
+    # Issue #38
+    if sys.version_info.minor >= 9:
+      self.assertEqual(monitor_dict, monitor_dict | reference_dict)
+    else:
+      self.assertEqual(monitor_dict, {**monitor_dict, **reference_dict})
 
   @mock.patch('time.time')
   def test_early_exit(self, mock_time):

compiler_opt/rl/data_reader.py

@@ -165,16 +165,13 @@ def create_sequence_example_dataset_fn(
   def _sequence_example_dataset_fn(sequence_examples):
     # Data collector returns empty strings for corner cases, filter them out
     # here.
-    dataset = (tf.data.Dataset.from_tensor_slices(sequence_examples)
-              .filter(lambda string: tf.strings.length(string) > 0)
-              .map(parser_fn)
-              .filter(lambda traj: tf.size(traj.reward) > 2)
-              .unbatch()
-              .batch(train_sequence_length, drop_remainder=True)
-              .cache()
-              .shuffle(trajectory_shuffle_buffer_size)
-              .batch(batch_size, drop_remainder=True)
-              )
+    dataset = (
+        tf.data.Dataset.from_tensor_slices(sequence_examples).filter(
+            lambda string: tf.strings.length(string) > 0).map(parser_fn).filter(
+                lambda traj: tf.size(traj.reward) > 2).unbatch().batch(
+                    train_sequence_length, drop_remainder=True).cache().shuffle(
+                        trajectory_shuffle_buffer_size).batch(
+                            batch_size, drop_remainder=True))
     return dataset
 
   return _sequence_example_dataset_fn

compiler_opt/rl/local_data_collector_test.py

@@ -19,6 +19,7 @@
 import multiprocessing as mp
 import string
 import subprocess
+import sys
 from unittest import mock
 
 import tensorflow as tf
@@ -147,7 +148,15 @@ def _test_iterator_fn(data_list):
             'total_trajectory_length': 18,
         }
     }
-    self.assertEqual(monitor_dict, monitor_dict | expected_monitor_dict_subset)
+    # Issue #38
+    if sys.version_info.minor >= 9:
+      self.assertEqual(monitor_dict,
+                       monitor_dict | expected_monitor_dict_subset)
+    else:
+      self.assertEqual(monitor_dict, {
+          **monitor_dict,
+          **expected_monitor_dict_subset
+      })
 
     data_iterator, monitor_dict = collector.collect_data(policy_path='policy')
     data = list(data_iterator)
@@ -158,7 +167,15 @@ def _test_iterator_fn(data_list):
             'total_trajectory_length': 18,
         }
     }
-    self.assertEqual(monitor_dict, monitor_dict | expected_monitor_dict_subset)
+    # Issue #38
+    if sys.version_info.minor >= 9:
+      self.assertEqual(monitor_dict,
+                       monitor_dict | expected_monitor_dict_subset)
+    else:
+      self.assertEqual(monitor_dict, {
+          **monitor_dict,
+          **expected_monitor_dict_subset
+      })
 
     collector.close_pool()