code: def testDataContext(self):
      agent = MyAgent(
          training_data_spec=(
              trajectory.Trajectory(
                  observation={"obs": tf.TensorSpec([], tf.float32)},
                  action=(),
                  policy_info={"info": tf.TensorSpec([], tf.int32)},
                  reward=tf.TensorSpec([], tf.float32, name="reward"),
                  step_type=tf.TensorSpec(
                      [], tf.int32, name="step_type"
                  ),
                  next_step_type=tf.TensorSpec(
                      [], tf.int32, name="next_step_type"
                  ),
                  discount=tensor_spec.BoundedTensorSpec(
                      [], tf.float32, 0.0, 1.0, name="discount"
                  ),
              )
          )
      )
      self.assertEqual(
          agent.data_context.time_step_spec,
          ts.time_step_spec({"obs": tf.TensorSpec([], tf.float32)}),
      )
      self.assertEqual(
          agent.collect_data_context.time_step_spec,
          ts.time_step_spec({"obs": tf.TensorSpec([], tf.float32)}),
      )
      self.assertEqual(agent.data_context.action_spec, ())
      self.assertEqual(agent.collect_data_context.action_spec, ())
      self.assertEqual(
          agent.data_context.info_spec,
          {"info": tf.TensorSpec([], tf.int32)},
      )
      self.assertEqual(agent.collect_data_context.info_spec, ())
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import numpy
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import test_util
  from tf_agents.agents import tf_agent
  from tf_agents.policies import random_tf_policy
  from tf_agents.specs import array_spec
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import test_utils
label_desc: Use image node_id 2 for calling the TFAgentTest obj's underlying member method code with example usage: obj.testDataContext() without return types