Add check is_valid_upload_filename routine to SwitchModel

Author: keigohtr

drucker/drucker_dashboard_servicer.py

@@ -149,17 +149,20 @@ def SwitchModel(self,
                     ) -> drucker_pb2.ModelResponse:
         """ Switch your ML model to run.
         """
+        if not self.is_valid_upload_filename(request.path):
+            raise Exception(f'Error: Invalid model path specified -> {request.path}')
+
         model_assignment = self.app.db.session.query(ModelAssignment).filter(ModelAssignment.service_name == self.app.config.SERVICE_NAME).one()
         model_assignment.model_path = request.path
         model_assignment.first_boot = False
         self.app.db.session.commit()
-        model_path = self.app.get_model_path()
 
         # :TODO: Use enum for SERVICE_INFRA
         if self.app.config.SERVICE_INFRA == "kubernetes":
             pass
         elif self.app.config.SERVICE_INFRA == "default":
-            self.app.load_model(model_path)
+            self.app.model_path = self.app.get_model_path()
+            self.app.load_model()
 
         return drucker_pb2.ModelResponse(status=1,
                                          message='Success: Switching model file.')
@@ -174,7 +177,7 @@ def EvaluateModel(self,
         first_req = next(request_iterator)
         save_path = first_req.data_path
         if not self.is_valid_upload_filename(save_path):
-            raise Exception(f'Error: Invalid evaluation file specified -> {save_path}')
+            raise Exception(f'Error: Invalid evaluation file path specified -> {save_path}')
 
         test_data = b''.join([first_req.data] + [r.data for r in request_iterator])
         result, details = self.app.evaluate(test_data)

drucker/test/test_dashboard_servicer.py

@@ -15,7 +15,7 @@ class DruckerWorkerServicerTest(unittest.TestCase):
     def test_ServiceInfo(self):
         servicer = DruckerDashboardServicer(logger=system_logger, app=app)
         request = drucker_pb2.ServiceInfoRequest()
-        response = servicer.ServiceInfo(request=request, context=Mock())
+        response = servicer.ServiceInfo(request, Mock())
         self.assertEqual(response.application_name, 'test')
         self.assertEqual(response.service_name, 'test-001')
         self.assertEqual(response.service_level, 'development')
@@ -73,7 +73,22 @@ def test_SwitchModel(self, mock_app):
         response = servicer.SwitchModel(request, Mock())
 
         self.assertEqual(response.status, 1)
-        mock_app.load_model.assert_called_once_with('test/my_path')
+        mock_app.load_model.assert_called_once_with()
+
+    @patch('drucker.test.DummyApp')
+    @patch('drucker.drucker_dashboard_servicer.Path')
+    def test_InvalidSwitchModel(self, mock_path_class, mock_app):
+        # mock setting
+        mock_path_class.return_value = Mock()
+        mock_path_class.return_value.name = 'my_path'
+        mock_app.get_model_path.return_value = 'test/my_path'
+        mock_app.config.SERVICE_INFRA = 'default'
+
+        servicer = DruckerDashboardServicer(logger=system_logger, app=mock_app)
+        request = drucker_pb2.SwitchModelRequest(path='../../my_path')
+        response = servicer.SwitchModel(request, Mock())
+
+        self.assertEqual(response.status, 0)
 
     @patch("builtins.open", new_callable=mock_open)
     @patch('drucker.drucker_dashboard_servicer.pickle')