fix tests

beam-cloud · Jan 11, 2024 · ac746fc · ac746fc
1 parent 025f384
commit ac746fc
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/sdk/tests/test_function.py b/sdk/tests/test_function.py
@@ -26,15 +26,19 @@ class TestTaskQueue(TestCase):
     def test_init(self):
         mock_stub = MagicMock()
 
-        queue = Function(Image(python_version="python3.8"), cpu=100, memory=128)
+        queue = Function(
+            cpu=100,
+            memory=128,
+            image=Image(python_version="python3.8"),
+        )
         queue.stub = mock_stub
 
         self.assertEqual(queue.image.python_version, "python3.8")
         self.assertEqual(queue.cpu, 100)
         self.assertEqual(queue.memory, 128)
 
     def test_run_local(self):
-        @Function(Image(python_version="python3.8"), cpu=100, memory=128)
+        @Function(cpu=100, memory=128, image=Image(python_version="python3.8"))
         def test_func():
             return 1
 
@@ -43,7 +47,7 @@ def test_func():
         self.assertEqual(resp, 1)
 
     def test_function_invoke(self):
-        @Function(Image(python_version="python3.8"), cpu=100, memory=128)
+        @Function(cpu=100, memory=128, image=Image(python_version="python3.8"))
         def test_func(*args, **kwargs):
             return 1998
 
@@ -67,7 +71,7 @@ def test_func(*args, **kwargs):
         self.assertRaises(SystemExit, test_func)
 
     def test_map(self):
-        @Function(Image(python_version="python3.8"), cpu=100, memory=128)
+        @Function(cpu=100, memory=128, image=Image(python_version="python3.8"))
         def test_func(*args, **kwargs):
             return 1998
 

diff --git a/sdk/tests/test_task_queue.py b/sdk/tests/test_task_queue.py
@@ -13,15 +13,15 @@ class TestTaskQueue(TestCase):
     def test_init(self):
         mock_stub = MagicMock()
 
-        queue = TaskQueue(Image(python_version="python3.8"), cpu=100, memory=128)
+        queue = TaskQueue(cpu=100, memory=128, image=Image(python_version="python3.8"))
         queue.stub = mock_stub
 
         self.assertEqual(queue.image.python_version, "python3.8")
         self.assertEqual(queue.cpu, 100)
         self.assertEqual(queue.memory, 128)
 
     def test_run_local(self):
-        @TaskQueue(Image(python_version="python3.8"), cpu=100, memory=128)
+        @TaskQueue(cpu=100, memory=128, image=Image(python_version="python3.8"))
         def test_func():
             return 1
 
@@ -30,7 +30,7 @@ def test_func():
         self.assertEqual(resp, 1)
 
     def test_put(self):
-        @TaskQueue(Image(python_version="python3.8"), cpu=100, memory=128)
+        @TaskQueue(cpu=100, memory=128, image=Image(python_version="python3.8"))
         def test_func():
             return 1
 
@@ -49,7 +49,7 @@ def test_func():
         self.assertRaises(SystemExit, test_func.put)
 
     def test__call__(self):
-        @TaskQueue(Image(python_version="python3.8"), cpu=100, memory=128)
+        @TaskQueue(cpu=100, memory=128, image=Image(python_version="python3.8"))
         def test_func():
             return 1