code: def test_correct_piece(self, batch_size):
      env = pbe.PiecewiseBernoulliPyEnvironment(
          [[0.1, 0.2, 0.3], [0.3, 0.2, 0.1], [0.1, 0.12, 0.14]],
          self.deterministic_duration_generator(),
          batch_size,
      )
      for t in range(100):
          env.reset()
          self.assertEqual(int(t / 10) % 3, env._current_piece)
          _ = env.step([0])
  
file_path: agents/tf_agents/bandits/environments/piecewise_bernoulli_py_environment_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import tensorflow
  from tf_agents.bandits.environments import piecewise_bernoulli_py_environment
label_desc: Use image node_id 3 for calling the PiecewiseBernoulliBanditPyEnvironmentTest obj's underlying member method code with example usage: obj.test_correct_piece(batch_size) without return types