code: def test_after_train_step_fn_with_stale_data(
      self, create_strategy_fn
  ):
      strategy = create_strategy_fn()
      with strategy.scope():
          # Prepare the test context context.
          train_step = train_utils.create_train_step()
          train_step.assign(225)
          train_steps_per_policy_update = 100
  
          # Create the after train function to test, and the test input.
          after_train_step_fn = train_utils.create_staleness_metrics_after_train_step_fn(
              train_step,
              train_steps_per_policy_update=train_steps_per_policy_update,
          )
          observation_train_steps = np.array(
              [[100], [200]], dtype=np.int64
          )
  
          # Define the expectations (expected scalar summary calls).
          expected_scalar_summary_calls = [
              mock.call(
                  name="staleness/max_train_step_delta_in_batch",
                  data=100,
                  step=225,
              ),
              mock.call(
                  name="staleness/max_policy_update_delta_in_batch",
                  data=1,
                  step=225,
              ),
              mock.call(
                  name="staleness/num_stale_obserations_in_batch",
                  data=1,
                  step=225,
              ),
          ]
  
          # Call the after train function and check the expectations.
          with mock.patch.object(
              tf.summary, "scalar", new=_SUMMARY_MOCK
          ), mock.patch.object(
              tf.compat.v2.summary, "scalar", new=_SUMMARY_MOCK
          ):
              # Call the `after_train_function` on the test input. Assumed the
              # observation train steps are stored in the field `priority` of the
              # the sample info of Reverb.
              info = reverb.SampleInfo(
                  *[None for _ in reverb.SampleInfo.tf_dtypes()]
              )
              info = info._replace(priority=observation_train_steps)
              strategy.run(
                  after_train_step_fn, args=((None, info), None)
              )
  
              # Check if the expected calls happened on the scalar summary.
              _SUMMARY_MOCK.assert_has_calls(
                  expected_scalar_summary_calls, any_order=False
              )
  
file_path: agents/tf_agents/train/utils/train_utils_test.py
imports: from absl.testing import parameterized
  from absl.testing.absltest import mock
  import numpy
  import reverb
  import tensorflow
  from tf_agents.system import system_multiprocessing
  from tf_agents.train.utils import test_utils
  from tf_agents.train.utils import train_utils
  from tf_agents.utils import test_utils
label_desc: Use image node_id 3 for calling the TrainUtilsTest obj's underlying member method code with example usage: obj.test_after_train_step_fn_with_stale_data(create_strategy_fn) without return types