def test_get_analysis_dataset_keys(self, preprocessing_fn, full_dataset_keys, cached_dataset_keys, expected_dataset_keys, expected_flat_data_required): # We force all dataset keys with entries in the cache dict will have a cache # hit. mocked_cache_entry_key = b'M' input_cache = { key: { mocked_cache_entry_key: 'C' } for key in cached_dataset_keys } feature_spec = {'x': tf.io.FixedLenFeature([], tf.float32)} with mock.patch( 'tensorflow_transform.beam.analysis_graph_builder.' 'analyzer_cache.make_cache_entry_key', return_value=mocked_cache_entry_key): dataset_keys, flat_data_required = ( analysis_graph_builder.get_analysis_dataset_keys( preprocessing_fn, feature_spec, full_dataset_keys, input_cache)) dot_string = nodes.get_dot_graph( [analysis_graph_builder._ANALYSIS_GRAPH]).to_string() self.WriteRenderedDotFile(dot_string) self.assertCountEqual(expected_dataset_keys, dataset_keys) self.assertEqual(expected_flat_data_required, flat_data_required)
def test_get_analysis_dataset_keys(self, preprocessing_fn, full_dataset_keys, cached_dataset_keys, expected_dataset_keys, use_tf_compat_v1): if not use_tf_compat_v1: test_case.skip_if_not_tf2('Tensorflow 2.x required') full_dataset_keys = [ analysis_graph_builder.analyzer_cache.DatasetKey(k) for k in full_dataset_keys ] # We force all dataset keys with entries in the cache dict will have a cache # hit. mocked_cache_entry_key = b'M' input_cache = { key: { mocked_cache_entry_key: 'C' } for key in cached_dataset_keys } feature_spec = {'x': tf.io.FixedLenFeature([], tf.float32)} specs = (feature_spec if use_tf_compat_v1 else impl_helper.get_type_specs_from_feature_specs(feature_spec)) with mock.patch( 'tensorflow_transform.beam.analysis_graph_builder.' 'analyzer_cache.make_cache_entry_key', return_value=mocked_cache_entry_key): dataset_keys = (analysis_graph_builder.get_analysis_dataset_keys( preprocessing_fn, specs, full_dataset_keys, input_cache, force_tf_compat_v1=use_tf_compat_v1)) dot_string = nodes.get_dot_graph( [analysis_graph_builder._ANALYSIS_GRAPH]).to_string() self.WriteRenderedDotFile(dot_string) self.assertCountEqual(expected_dataset_keys, dataset_keys)