[TGA] broken test

tests/test_tga.py

@@ -129,25 +129,31 @@ def test_full_context_block(self):
         obs = {'text': '1 2 3 4 ' * 256, 'episode_done': False}
         agent.observe(obs)
         batch = agent.batchify([agent.observation])
-        self.assertEqual(agent._get_context(batch, 0).tolist(), [5, 4, 6, 7] * 256)
+        self.assertEqual(
+            agent._get_batch_context(batch)[0].tolist(), [5, 4, 6, 7] * 256
+        )
 
         # observe 1 more obs, context is the same (truncation)
         agent.observe(obs)
         batch = agent.batchify([agent.observation])
-        self.assertEqual(agent._get_context(batch, 0).tolist(), [5, 4, 6, 7] * 256)
+        self.assertEqual(
+            agent._get_batch_context(batch)[0].tolist(), [5, 4, 6, 7] * 256
+        )
 
         # Now, set agent's beam_block_full_context
         args += ['--beam-block-full-context', 'true']
         agent2 = create_agent(pp.parse_args(args), True)
         agent2.observe(obs)
         batch = agent2.batchify([agent2.observation])
-        self.assertEqual(agent2._get_context(batch, 0).tolist(), [5, 4, 6, 7] * 256)
+        self.assertEqual(
+            agent2._get_batch_context(batch)[0].tolist(), [5, 4, 6, 7] * 256
+        )
 
         # observe 1 more obs, context is larger now
         agent2.observe(obs)
         batch = agent2.batchify([agent2.observation])
         self.assertEqual(
-            agent2._get_context(batch, 0).tolist(),
+            agent2._get_batch_context(batch)[0].tolist(),
             [5, 4, 6, 7] * 256 + [3] + [5, 4, 6, 7] * 256,
         )  # 3 is end token.