code: def assertTfDataSetExamplesMatchExpected(
      self,
      input_fn,
      expected_dataset,
      use_tpu=False,
      workaround_list_files=False,
  ):
      # Note that we use input_fn to get an iterator, while we use
      # expected_dataset to get a filename, even though they are the same
      # type (DeepVariantInput), and may even be the same object.
      with tf.compat.v1.Session() as sess:
          params = {"batch_size": 1}
          batch_feed = tf.compat.v1.data.make_one_shot_iterator(
              input_fn(params)
          ).get_next()
  
          sess.run(tf.compat.v1.global_variables_initializer())
          sess.run(tf.compat.v1.local_variables_initializer())
          seen = []
          while True:
              try:
                  features, _ = sess.run(batch_feed)
              except tf.errors.OutOfRangeError:
                  break
              locus = features["locus"][0]
              if use_tpu:
                  locus = dv_utils.int_tensor_to_string(locus)
              # NB, this looks like: array(['chr20:10001019-10001019'], dtype=object)
              seen.append(locus)
  
      if workaround_list_files:
          # This really only works for loci, because those are string valued and
          # are expected to show up in sorted order.  For arbitrary data that's
          # not true.  In prod we have the version of tf that lets us turn off
          # shuffling so this path is skipped, but kokoro hits this.
          seen = sorted(seen)
  
      expected_loci = [
          example.features.feature["locus"].bytes_list.value[0]
          for example in tfrecord.read_tfrecords(
              expected_dataset.input_file_spec
          )
      ]
      self.assertLen(expected_loci, expected_dataset.num_examples)
      if seen != expected_loci:
          print("\n\nlen expected seen", len(expected_loci), len(seen))
          print("\n\nexpected=", expected_loci)
          print("\n\nseen=", seen)
      self.assertEqual(expected_loci, seen)
      # Note that this expected shape comes from the golden dataset. If the data
      # is remade in the future, the values might need to be modified accordingly.
      self.assertEqual(
          dv_constants.PILEUP_DEFAULT_DIMS,
          expected_dataset.tensor_shape,
      )
  
imports: import math
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tensorflow import estimator
  from deepvariant import data_providers
  from deepvariant import dv_config
  from deepvariant import dv_constants
  from deepvariant import dv_utils
  from deepvariant import testdata
  from deepvariant.protos import deepvariant_pb2
  from third_party.nucleus.io import tfrecord
  from third_party.nucleus.testing import test_utils
  from third_party.nucleus.util import variant_utils
  from tensorflow.core.example import example_pb2
label_desc: Use image node_id 6 for calling the DataProviderTest obj's underlying member method code with example usage: obj.assertTfDataSetExamplesMatchExpected(input_fn, expected_dataset, use_tpu, workaround_list_files) without return types