docstr: Test that dataset is being shuffled when asked.
docstr_tok: ["Test","that","dataset","is","being","shuffled","when","asked","."]
code: def testShuffle(self):
      """Test that dataset is being shuffled when asked."""
      # 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))
          .repeat()
          .shuffle(4 * batch_size)
      )
      reward_distribution = deterministic_reward_distribution(tf.eye(3))
  
      # Note - shuffle should hapen *first* in call chain, so this
      # test will fail if shuffle is called e.g. after batch or prefetch.
      dataset.shuffle = mock.Mock(
          spec=dataset.shuffle, side_effect=dataset.shuffle
      )
      ce.ClassificationBanditEnvironment(
          dataset, reward_distribution, batch_size
      )
      dataset.shuffle.assert_not_called()
      ce.ClassificationBanditEnvironment(
          dataset,
          reward_distribution,
          batch_size,
          shuffle_buffer_size=3,
          seed=7,
      )
      dataset.shuffle.assert_called_with(
          buffer_size=3, reshuffle_each_iteration=True, seed=7
      )
  
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 5 for calling the ClassificationEnvironmentTest obj's underlying member method code with example usage: obj.testShuffle() without return types