text: import unittest
  
  from d2go.export.torchscript import (
      MobileOptimizationConfig,
      update_export_kwargs_from_export_method,
  )
  
  
  @update_export_kwargs_from_export_method
  def mock_export(cls, model, input_args, save_path, export_method, **export_kwargs):
      # Return the export kwargs, so that we can check to make sure it's set as expected
      return export_kwargs
  
  
  class TestTorchscriptExportMethods(unittest.TestCase):
      def test_update_export_kwargs_from_export_method(self):
  
          _empty_export_kwargs = {}
  
          def try_mock_export(export_method: str, export_kwargs=_empty_export_kwargs):
              return mock_export(
                  cls=None,
                  model=None,
                  input_args=None,
                  save_path=None,
                  export_method=export_method,
                  **export_kwargs
              )
  
          export_method_string = "torchscript"
          new_export_kwargs = try_mock_export(export_method_string)
          self.assertNotIn("mobile_optimization", new_export_kwargs)
  
          export_method_string = "torchscript_mobile"
          new_export_kwargs = try_mock_export(export_method_string)
          self.assertIn("mobile_optimization", new_export_kwargs)
          self.assertEquals(
              type(new_export_kwargs["mobile_optimization"]),
              MobileOptimizationConfig,
          )
          self.assertEquals(new_export_kwargs["mobile_optimization"].backend, "CPU")
  
          export_method_string = "torchscript_mobile-metal"
          new_export_kwargs = try_mock_export(export_method_string)
          self.assertEquals(new_export_kwargs["mobile_optimization"].backend, "metal")
  
          export_method_string = "torchscript_mobile-vulkan"
          new_export_kwargs = try_mock_export(export_method_string)
          self.assertEquals(new_export_kwargs["mobile_optimization"].backend, "vulkan")
  
          export_method_string = "torchscript_mobile@tracing"
          new_export_kwargs = try_mock_export(export_method_string)
          self.assertEquals(new_export_kwargs["jit_mode"], "trace")
  
          export_method_string = "torchscript_mobile@scripting"
          new_export_kwargs = try_mock_export(export_method_string)
          self.assertEquals(new_export_kwargs["jit_mode"], "script")
  