fix tests

beam-cloud · Jan 11, 2024 · b612a33 · b612a33
1 parent 47b3537
commit b612a33
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 11 deletions.
diff --git a/sdk/Makefile b/sdk/Makefile
@@ -7,7 +7,7 @@ docs:
 	@poetry run ./bin/docgen.py 2>/dev/null
 
 tests:
-	BEAM_SERIALIZE_MODE=start poetry run pytest -s
+	poetry run pytest -s
 
 format:
 	poetry run black .

diff --git a/sdk/tests/test_queue.py b/sdk/tests/test_queue.py
@@ -22,14 +22,14 @@ def setUp(self):
     def test_put(self):
         mock_stub = MagicMock()
 
-        mock_stub.put = mock_coroutine_with_result(SimpleQueuePutResponse(ok=True))
+        mock_stub.simple_queue_put = mock_coroutine_with_result(SimpleQueuePutResponse(ok=True))
 
         queue = SimpleQueue(name="test")
         queue.stub = mock_stub
 
         self.assertTrue(queue.put("test"))
 
-        mock_stub.put = mock_coroutine_with_result(SimpleQueuePutResponse(ok=False))
+        mock_stub.simple_queue_put = mock_coroutine_with_result(SimpleQueuePutResponse(ok=False))
 
         queue = SimpleQueue(name="test")
         queue.stub = mock_stub
@@ -41,7 +41,7 @@ def test_pop(self):
 
         pickled_value = cloudpickle.dumps("test")
 
-        mock_stub.pop = mock_coroutine_with_result(
+        mock_stub.simple_queue_pop = mock_coroutine_with_result(
             SimpleQueuePopResponse(ok=True, value=pickled_value)
         )
 
@@ -50,7 +50,9 @@ def test_pop(self):
 
         self.assertEqual(queue.pop(), "test")
 
-        mock_stub.pop = mock_coroutine_with_result(SimpleQueuePopResponse(ok=False, value=b""))
+        mock_stub.simple_queue_pop = mock_coroutine_with_result(
+            SimpleQueuePopResponse(ok=False, value=b"")
+        )
 
         queue = SimpleQueue(name="test")
         queue.stub = mock_stub
@@ -62,7 +64,7 @@ def test_peek(self):
 
         pickled_value = cloudpickle.dumps("test")
 
-        mock_stub.peek = mock_coroutine_with_result(
+        mock_stub.simple_queue_peek = mock_coroutine_with_result(
             SimpleQueuePeekResponse(ok=True, value=pickled_value)
         )
 
@@ -71,7 +73,9 @@ def test_peek(self):
 
         self.assertEqual(queue.peek(), "test")
 
-        mock_stub.peek = mock_coroutine_with_result(SimpleQueuePeekResponse(ok=False, value=b""))
+        mock_stub.simple_queue_peek = mock_coroutine_with_result(
+            SimpleQueuePeekResponse(ok=False, value=b"")
+        )
 
         queue = SimpleQueue(name="test")
         queue.stub = mock_stub
@@ -81,14 +85,18 @@ def test_peek(self):
     def test_empty(self):
         mock_stub = MagicMock()
 
-        mock_stub.empty = mock_coroutine_with_result(SimpleQueueEmptyResponse(ok=True, empty=True))
+        mock_stub.simple_queue_empty = mock_coroutine_with_result(
+            SimpleQueueEmptyResponse(ok=True, empty=True)
+        )
 
         queue = SimpleQueue(name="test")
         queue.stub = mock_stub
 
         self.assertTrue(queue.empty())
 
-        mock_stub.empty = mock_coroutine_with_result(SimpleQueueEmptyResponse(ok=False, empty=True))
+        mock_stub.simple_queue_empty = mock_coroutine_with_result(
+            SimpleQueueEmptyResponse(ok=False, empty=True)
+        )
 
         queue = SimpleQueue(name="test")
         queue.stub = mock_stub
@@ -98,14 +106,18 @@ def test_empty(self):
     def test_size(self):
         mock_stub = MagicMock()
 
-        mock_stub.size = mock_coroutine_with_result(SimpleQueueSizeResponse(ok=True, size=1))
+        mock_stub.simple_queue_size = mock_coroutine_with_result(
+            SimpleQueueSizeResponse(ok=True, size=1)
+        )
 
         queue = SimpleQueue(name="test")
         queue.stub = mock_stub
 
         self.assertEqual(len(queue), 1)
 
-        mock_stub.size = mock_coroutine_with_result(SimpleQueueSizeResponse(ok=False, size=1))
+        mock_stub.simple_queue_size = mock_coroutine_with_result(
+            SimpleQueueSizeResponse(ok=False, size=1)
+        )
 
         queue = SimpleQueue(name="test")
         queue.stub = mock_stub