code: def testAction(self):
      py_observation_spec = array_spec.BoundedArraySpec(
          (3,), np.int32, 1, 1
      )
      py_time_step_spec = ts.time_step_spec(py_observation_spec)
      py_action_spec = array_spec.BoundedArraySpec((7,), np.int32, 1, 1)
      py_policy_state_spec = array_spec.BoundedArraySpec(
          (5,), np.int32, 0, 1
      )
      py_policy_info_spec = array_spec.BoundedArraySpec(
          (3,), np.int32, 0, 1
      )
  
      mock_py_policy = mock.create_autospec(py_policy.PyPolicy)
      mock_py_policy.time_step_spec = py_time_step_spec
      mock_py_policy.action_spec = py_action_spec
      mock_py_policy.policy_state_spec = py_policy_state_spec
      mock_py_policy.info_spec = py_policy_info_spec
  
      expected_py_policy_state = np.ones(
          py_policy_state_spec.shape, py_policy_state_spec.dtype
      )
      expected_py_time_step = tf.nest.map_structure(
          lambda arr_spec: np.ones(
              (1,) + arr_spec.shape, arr_spec.dtype
          ),
          py_time_step_spec,
      )
      expected_py_action = np.ones(
          (1,) + py_action_spec.shape, py_action_spec.dtype
      )
      expected_new_py_policy_state = np.zeros(
          py_policy_state_spec.shape, py_policy_state_spec.dtype
      )
      expected_py_info = np.zeros(
          py_policy_info_spec.shape, py_policy_info_spec.dtype
      )
  
      mock_py_policy.action.return_value = policy_step.PolicyStep(
          nest_utils.unbatch_nested_array(expected_py_action),
          expected_new_py_policy_state,
          expected_py_info,
      )
  
      tf_mock_py_policy = tf_py_policy.TFPyPolicy(mock_py_policy)
      time_step = tf.nest.map_structure(
          lambda arr_spec: tf.ones(
              (1,) + arr_spec.shape, arr_spec.dtype
          ),
          py_time_step_spec,
      )
      action_step = tf_mock_py_policy.action(
          time_step, tf.ones(py_policy_state_spec.shape, tf.int32)
      )
      py_action_step = self.evaluate(action_step)
  
      self.assertEqual(1, mock_py_policy.action.call_count)
      np.testing.assert_equal(
          mock_py_policy.action.call_args[1]["time_step"],
          nest_utils.unbatch_nested_array(expected_py_time_step),
      )
      np.testing.assert_equal(
          mock_py_policy.action.call_args[1]["policy_state"],
          expected_py_policy_state,
      )
      np.testing.assert_equal(py_action_step.action, expected_py_action)
      np.testing.assert_equal(
          py_action_step.state, expected_new_py_policy_state
      )
      np.testing.assert_equal(py_action_step.info, expected_py_info)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing.absltest import mock
  import numpy
  import tensorflow
  from tf_agents.policies import py_policy
  from tf_agents.policies import random_py_policy
  from tf_agents.policies import tf_py_policy
  from tf_agents.specs import array_spec
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.utils import nest_utils
  from tf_agents.utils import test_utils
label_desc: Use image node_id 2 for calling the TFPyPolicyTest obj's underlying member method code with example usage: obj.testAction() without return types