code: def testMaybeAddOneAction(self, batch_size, num_actions):
      original_mask = tf.eye(batch_size, num_actions, dtype=tf.int32)
      new_mask = self.evaluate(
          masked_tf_env._maybe_add_one_action(original_mask)
      )
      self.assertAllEqual(
          original_mask[:num_actions, :], new_mask[:num_actions, :]
      )
      ones = tf.ones([batch_size], dtype=tf.int32)
      self.assertAllEqual(tf.reduce_max(new_mask, axis=1), ones)
      self.assertAllEqual(tf.reduce_sum(new_mask, axis=1), ones)
  
file_path: agents/tf_agents/bandits/environments/bernoulli_action_mask_tf_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
  import tensorflow_probability
  from tf_agents.bandits.environments import bernoulli_action_mask_tf_environment
  from tf_agents.bandits.environments import random_bandit_environment
  from tf_agents.specs import tensor_spec
label_desc: Use image node_id 1 for calling the BernoulliActionMaskTfEnvironmentTest obj's underlying member method code with example usage: obj.testMaybeAddOneAction(batch_size, num_actions) without return types