Code health update on model server tests

lit_nlp/examples/gcp/model_server_test.py

@@ -2,21 +2,23 @@
 from unittest import mock
 
 from absl.testing import absltest
+from absl.testing import parameterized
 from lit_nlp.examples.gcp import model_server
 from lit_nlp.examples.prompt_debugging import utils as pd_utils
 import webtest
 
 
-class TestWSGIApp(absltest.TestCase):
+class TestWSGIApp(parameterized.TestCase):
 
-  @mock.patch('lit_nlp.examples.prompt_debugging.models.get_models')
-  def test_predict_endpoint(self, mock_get_models):
+  @classmethod
+  def setUpClass(cls):
     test_model_name = 'lit_on_gcp_test_model'
+    sal_name, tok_name = pd_utils.generate_model_group_names(test_model_name)
     test_model_config = f'{test_model_name}:test_model_path'
     os.environ['MODEL_CONFIG'] = test_model_config
 
-    mock_model = mock.MagicMock()
-    mock_model.predict.side_effect = [[{'response': 'test output text'}]]
+    generation_model = mock.MagicMock()
+    generation_model.predict.side_effect = [[{'response': 'test output text'}]]
 
     salience_model = mock.MagicMock()
     salience_model.predict.side_effect = [[{
@@ -30,33 +32,42 @@ def test_predict_endpoint(self, mock_get_models):
         [{'tokens': ['test', 'output', 'text']}]
     ]
 
-    sal_name, tok_name = pd_utils.generate_model_group_names(test_model_name)
-
-    mock_get_models.return_value = {
-        test_model_name: mock_model,
+    cls.mock_models = {
+        test_model_name: generation_model,
         sal_name: salience_model,
         tok_name: tokenize_model,
     }
-    app = webtest.TestApp(model_server.get_wsgi_app())
 
-    response = app.post_json('/predict', {'inputs': 'test_input'})
-    self.assertEqual(response.status_code, 200)
-    self.assertEqual(response.json, [{'response': 'test output text'}])
 
-    response = app.post_json('/salience', {'inputs': 'test_input'})
-    self.assertEqual(response.status_code, 200)
-    self.assertEqual(
-        response.json,
-        [{
+  @parameterized.named_parameters(
+    dict(
+        testcase_name='predict',
+        endpoint='/predict',
+        expected=[{'response': 'test output text'}],
+    ),
+    dict(
+        testcase_name='salience',
+        endpoint='/salience',
+        expected=[{
             'tokens': ['test', 'output', 'text'],
             'grad_l2': [0.1234, 0.3456, 0.5678],
             'grad_dot_input': [0.1234, -0.3456, 0.5678],
         }],
-    )
+    ),
+    dict(
+        testcase_name='tokenize',
+        endpoint='/tokenize',
+        expected=[{'tokens': ['test', 'output', 'text']}],
+    ),
+  )
+  @mock.patch('lit_nlp.examples.prompt_debugging.models.get_models')
+  def test_endpoint(self, mock_get_models, endpoint, expected):
+    mock_get_models.return_value = self.mock_models
+    app = webtest.TestApp(model_server.get_wsgi_app())
 
-    response = app.post_json('/tokenize', {'inputs': 'test_input'})
+    response = app.post_json(endpoint, {'inputs': [{'prompt': 'test input'}]})
     self.assertEqual(response.status_code, 200)
-    self.assertEqual(response.json, [{'tokens': ['test', 'output', 'text']}])
+    self.assertEqual(response.json, expected)
 
 
 if __name__ == '__main__':