docstr: Test that dataset is being prefetched when asked.
docstr_tok: ["Test","that","dataset","is","being","prefetched","when","asked","."]
code: def testPrefetch(self, mock_dataset_iterator):
      """Test that dataset is being prefetched when asked."""
      mock_dataset_iterator.return_value = "mock_iterator_result"
      # Reward of 1 is given if action == (context % 3)
      context = tf.reshape(tf.range(128), shape=[128, 1])
      labels = tf.math.mod(context, 3)
      batch_size = 32
      dataset = tf.data.Dataset.from_tensor_slices((context, labels))
      reward_distribution = deterministic_reward_distribution(tf.eye(3))
  
      # Operation order should be batch() then prefetch(), have to jump
      # through a couple hoops to get this sequence tested correctly.
  
      # Save dataset.prefetch in temp mock_prefetch, return batched dataset to
      # make down-stream logic work correctly with batch dimensions.
      batched_dataset = dataset.batch(batch_size)
      mock_prefetch = mock.Mock(
          spec=dataset.prefetch, return_value=batched_dataset
      )
      # Replace dataset.batch with mock batch that returns original dataset,
      # in order to make mocking out it's prefetch call easier.
      dataset.batch = mock.Mock(
          spec=batched_dataset, return_value=batched_dataset
      )
      # Replace dataset.prefetch with mock_prefetch.
      batched_dataset.prefetch = mock_prefetch
      env = ce.ClassificationBanditEnvironment(
          dataset, reward_distribution, batch_size, repeat_dataset=False
      )
      dataset.batch.assert_called_with(batch_size, drop_remainder=True)
      batched_dataset.prefetch.assert_not_called()
      mock_dataset_iterator.assert_called_with(batched_dataset)
      self.assertEqual(env._data_iterator, "mock_iterator_result")
      env = ce.ClassificationBanditEnvironment(
          dataset,
          reward_distribution,
          batch_size,
          repeat_dataset=False,
          prefetch_size=3,
      )
      dataset.batch.assert_called_with(batch_size, drop_remainder=True)
      batched_dataset.prefetch.assert_called_with(3)
      mock_dataset_iterator.assert_called_with(batched_dataset)
      self.assertEqual(env._data_iterator, "mock_iterator_result")
  
file_path: agents/tf_agents/bandits/environments/classification_environment_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  from absl.testing.absltest import mock
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.environments import classification_environment
label_desc: Use image node_id 6 for calling the ClassificationEnvironmentTest obj's underlying member method code with example usage: obj.testPrefetch(mock_dataset_iterator) without return types