1414
1515import os
1616import sys
17- import unittest
1817
1918
2019git_repo_path = os .path .abspath (os .path .dirname (os .path .dirname (os .path .dirname (__file__ ))))
2827check_dummies .PATH_TO_DIFFUSERS = os .path .join (git_repo_path , "src" , "diffusers" )
2928
3029
31- class CheckDummiesTester ( unittest . TestCase ) :
30+ class TestCheckDummies :
3231 def test_find_backend (self ):
3332 simple_backend = find_backend (" if not is_torch_available():" )
34- self . assertEqual ( simple_backend , "torch" )
33+ assert simple_backend == "torch"
3534
3635 # backend_with_underscore = find_backend(" if not is_tensorflow_text_available():")
37- # self.assertEqual( backend_with_underscore, "tensorflow_text")
36+ # assert backend_with_underscore == "tensorflow_text"
3837
3938 double_backend = find_backend (" if not (is_torch_available() and is_transformers_available()):" )
40- self . assertEqual ( double_backend , "torch_and_transformers" )
39+ assert double_backend == "torch_and_transformers"
4140
4241 # double_backend_with_underscore = find_backend(
4342 # " if not (is_sentencepiece_available() and is_tensorflow_text_available()):"
4443 # )
45- # self.assertEqual( double_backend_with_underscore, "sentencepiece_and_tensorflow_text")
44+ # assert double_backend_with_underscore == "sentencepiece_and_tensorflow_text"
4645
4746 triple_backend = find_backend (
4847 " if not (is_torch_available() and is_transformers_available() and is_onnx_available()):"
4948 )
50- self . assertEqual ( triple_backend , "torch_and_transformers_and_onnx" )
49+ assert triple_backend == "torch_and_transformers_and_onnx"
5150
5251 def test_read_init (self ):
5352 objects = read_init ()
5453 # We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
55- self . assertIn ( "torch" , objects )
56- self . assertIn ( "torch_and_transformers" , objects )
57- self . assertIn ( "torch_and_transformers_and_onnx" , objects )
54+ assert "torch" in objects
55+ assert "torch_and_transformers" in objects
56+ assert "torch_and_transformers_and_onnx" in objects
5857
5958 # Likewise, we can't assert on the exact content of a key
60- self . assertIn ( "UNet2DModel" , objects ["torch" ])
61- self . assertIn ( "StableDiffusionPipeline" , objects ["torch_and_transformers" ])
62- self . assertIn ( "LMSDiscreteScheduler" , objects ["torch_and_scipy" ])
63- self . assertIn ( "OnnxStableDiffusionPipeline" , objects ["torch_and_transformers_and_onnx" ])
59+ assert "UNet2DModel" in objects ["torch" ]
60+ assert "StableDiffusionPipeline" in objects ["torch_and_transformers" ]
61+ assert "LMSDiscreteScheduler" in objects ["torch_and_scipy" ]
62+ assert "OnnxStableDiffusionPipeline" in objects ["torch_and_transformers_and_onnx" ]
6463
6564 def test_create_dummy_object (self ):
6665 dummy_constant = create_dummy_object ("CONSTANT" , "'torch'" )
67- self . assertEqual ( dummy_constant , "\n CONSTANT = None\n " )
66+ assert dummy_constant == "\n CONSTANT = None\n "
6867
6968 dummy_function = create_dummy_object ("function" , "'torch'" )
70- self .assertEqual (
71- dummy_function , "\n def function(*args, **kwargs):\n requires_backends(function, 'torch')\n "
72- )
69+ assert dummy_function == "\n def function(*args, **kwargs):\n requires_backends(function, 'torch')\n "
7370
7471 expected_dummy_class = """
7572class FakeClass(metaclass=DummyObject):
@@ -87,7 +84,7 @@ def from_pretrained(cls, *args, **kwargs):
8784 requires_backends(cls, 'torch')
8885"""
8986 dummy_class = create_dummy_object ("FakeClass" , "'torch'" )
90- self . assertEqual ( dummy_class , expected_dummy_class )
87+ assert dummy_class == expected_dummy_class
9188
9289 def test_create_dummy_files (self ):
9390 expected_dummy_pytorch_file = """# This file is autogenerated by the command `make fix-copies`, do not edit.
@@ -116,4 +113,4 @@ def from_pretrained(cls, *args, **kwargs):
116113 requires_backends(cls, ["torch"])
117114"""
118115 dummy_files = create_dummy_files ({"torch" : ["CONSTANT" , "function" , "FakeClass" ]})
119- self . assertEqual ( dummy_files ["torch" ], expected_dummy_pytorch_file )
116+ assert dummy_files ["torch" ] == expected_dummy_pytorch_file
0 commit comments