fix translation server unittest

OpenNMT · Nov 7, 2019 · 285b3c7 · 285b3c7
1 parent 1979b34
commit 285b3c7
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/onmt/tests/test_translation_server.py b/onmt/tests/test_translation_server.py
@@ -112,13 +112,17 @@ def test_run(self):
         sm = ServerModel(opt, model_id, model_root=model_root, load=True)
         inp = [{"src": "hello how are you today"},
                {"src": "good morning to you ."}]
-        results, scores, n_best, time = sm.run(inp)
+        results, scores, n_best, time, aligns = sm.run(inp)
         self.assertIsInstance(results, list)
         for sentence_string in results:
             self.assertIsInstance(sentence_string, string_types)
         self.assertIsInstance(scores, list)
         for elem in scores:
             self.assertIsInstance(elem, float)
+        self.assertIsInstance(aligns, list)
+        for align_string in aligns:
+            if align_string is not None:
+                self.assertIsInstance(align_string, string_types)
         self.assertEqual(len(results), len(scores))
         self.assertEqual(len(scores), len(inp))
         self.assertEqual(n_best, 1)