Refactor tests to use Mocking for TPU sample operations

Author: Thoughtseize1

tpu/test_tpu.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from unittest.mock import MagicMock, patch
+
 import uuid
 
+from google.cloud.tpu_v2.services.tpu.pagers import ListNodesPager
 from google.cloud.tpu_v2.types import AcceleratorConfig, Node
 
 import pytest
@@ -28,68 +31,142 @@
 import stop_tpu
 
 
-TPU_NAME = "test-tpu-" + uuid.uuid4().hex[:10]
+TPU_NAME = "test-tpu-" + uuid.uuid4().hex[:6]
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 ZONE = "us-south1-a"
+FULL_TPU_NAME = f"projects/{PROJECT_ID}/locations/{ZONE}/nodes/{TPU_NAME}"
 TPU_TYPE = "v5litepod-1"
 TPU_VERSION = "tpu-vm-tf-2.17.0-pjrt"
+METADATA = {
+    "startup-script": """#!/bin/bash
+    echo "Hello World" > /var/log/hello.log
+    sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1
+    """
+}
+
+
+@pytest.fixture
+def mock_tpu_client() -> MagicMock:
+    with patch("google.cloud.tpu_v2.TpuClient") as mock_client:
+        yield mock_client.return_value
+
+
+@pytest.fixture
+def operation() -> MagicMock:
+    yield MagicMock()
+
+
+def test_creating_tpu(mock_tpu_client: MagicMock, operation: MagicMock) -> None:
+    mock_response = MagicMock(spec=Node)
+    mock_response.state = Node.State.READY
+    mock_response.name = FULL_TPU_NAME
+
+    mock_tpu_client.create_node.return_value = operation
+    operation.result.return_value = mock_response
+
+    tpu = create_tpu.create_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME, TPU_TYPE, TPU_VERSION)
+
+    assert tpu.name == FULL_TPU_NAME
+    assert tpu.state == Node.State.READY
+    mock_tpu_client.create_node.assert_called_once()
+    operation.result.assert_called_once()
+
+
+def test_delete_tpu(mock_tpu_client: MagicMock) -> None:
+    delete_tpu.delete_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME)
+    mock_tpu_client.delete_node.called_once()
 
 
-# Instance of TPU
-@pytest.fixture(scope="session")
-def tpu_instance() -> Node:
-    yield create_tpu.create_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME, TPU_TYPE, TPU_VERSION)
-    try:
-        delete_tpu.delete_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME)
-    except Exception as e:
-        print(f"Error during cleanup: {e}")
+def test_creating_with_startup_script(
+    mock_tpu_client: MagicMock, operation: MagicMock
+) -> None:
+    mock_response = MagicMock(spec=Node)
+    mock_response.metadata = METADATA
+    mock_tpu_client.create_node.return_value = operation
+    operation.result.return_value = mock_response
 
+    tpu_with_script = create_tpu_with_script.create_cloud_tpu_with_script(
+        PROJECT_ID, ZONE, TPU_NAME, TPU_TYPE, TPU_VERSION
+    )
 
-def test_creating_tpu(tpu_instance: Node) -> None:
-    assert tpu_instance.state == Node.State.READY
+    mock_tpu_client.create_node.assert_called_once()
+    operation.result.assert_called_once()
+    assert "--upgrade numpy" in tpu_with_script.metadata["startup-script"]
 
 
-def test_creating_with_startup_script() -> None:
-    tpu_name_with_script = "script-tpu-" + uuid.uuid4().hex[:5]
-    try:
-        tpu_with_script = create_tpu_with_script.create_cloud_tpu_with_script(
-            PROJECT_ID, ZONE, tpu_name_with_script, TPU_TYPE, TPU_VERSION
-        )
-        assert "--upgrade numpy" in tpu_with_script.metadata["startup-script"]
-    finally:
-        print(f"\n\n ------------ Deleting TPU {TPU_NAME}\n ------------")
-        delete_tpu.delete_cloud_tpu(PROJECT_ID, ZONE, tpu_name_with_script)
+def test_get_tpu(mock_tpu_client: MagicMock) -> None:
+    mock_response = MagicMock(spec=Node)
+    mock_response.name = FULL_TPU_NAME
+    mock_response.state = Node.State.READY
 
+    mock_tpu_client.get_node.return_value = mock_response
 
-def test_get_tpu() -> None:
     tpu = get_tpu.get_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME)
+
     assert tpu.state == Node.State.READY
-    assert tpu.name == f"projects/{PROJECT_ID}/locations/{ZONE}/nodes/{TPU_NAME}"
+    assert tpu.name == FULL_TPU_NAME
+    mock_tpu_client.get_node.assert_called_once()
+
+
+def test_list_tpu(mock_tpu_client: MagicMock) -> None:
+    mock_pager = MagicMock(spec=ListNodesPager)
+    nodes = [
+        Node(name="Node1", state=Node.State.READY),
+        Node(name="Node2", state=Node.State.CREATING),
+    ]
+    mock_pager.__iter__.return_value = nodes
 
+    mock_tpu_client.list_nodes.return_value = mock_pager
 
-def test_list_tpu() -> None:
     nodes = list_tpu.list_cloud_tpu(PROJECT_ID, ZONE)
     assert len(list(nodes)) > 0
+    mock_tpu_client.list_nodes.assert_called_once()
 
 
-def test_stop_tpu() -> None:
+def test_stop_tpu(mock_tpu_client: MagicMock, operation: MagicMock) -> None:
+    mock_response = MagicMock(spec=Node)
+    mock_response.state = Node.State.STOPPED
+
+    mock_tpu_client.stop_node.return_value = operation
+    operation.result.return_value = mock_response
+
     node = stop_tpu.stop_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME)
+
+    mock_tpu_client.stop_node.assert_called_once()
+    operation.result.assert_called_once()
     assert node.state == Node.State.STOPPED
 
 
-def test_start_tpu() -> None:
+def test_start_tpu(mock_tpu_client: MagicMock, operation: MagicMock) -> None:
+    mock_response = MagicMock(spec=Node)
+    mock_response.state = Node.State.READY
+
+    mock_tpu_client.start_node.return_value = operation
+    operation.result.return_value = mock_response
+
     node = start_tpu.start_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME)
+
+    mock_tpu_client.start_node.assert_called_once()
+    operation.result.assert_called_once()
     assert node.state == Node.State.READY
 
 
-def test_with_topology() -> None:
-    topology_tpu_name = "topology-tpu-" + uuid.uuid4().hex[:5]
-    topology_zone = "us-central1-a"
-    try:
-        topology_tpu = create_tpu_topology.create_cloud_tpu_with_topology(
-            PROJECT_ID, topology_zone, topology_tpu_name, TPU_VERSION
-        )
-        assert topology_tpu.accelerator_config.type_ == AcceleratorConfig.Type.V3
-        assert topology_tpu.accelerator_config.topology == "2x2"
-    finally:
-        delete_tpu.delete_cloud_tpu(PROJECT_ID, topology_zone, topology_tpu_name)
+def test_with_topology(mock_tpu_client: MagicMock, operation: MagicMock) -> None:
+    from google.cloud import tpu_v2
+
+    mock_response = MagicMock(spec=Node)
+    mock_response.accelerator_config = tpu_v2.AcceleratorConfig(
+        type_=tpu_v2.AcceleratorConfig.Type.V3,
+        topology="2x2",
+    )
+
+    mock_tpu_client.create_node.return_value = operation
+    operation.result.return_value = mock_response
+
+    topology_tpu = create_tpu_topology.create_cloud_tpu_with_topology(
+        PROJECT_ID, ZONE, TPU_NAME, TPU_VERSION
+    )
+    assert topology_tpu.accelerator_config.type_ == AcceleratorConfig.Type.V3
+    assert topology_tpu.accelerator_config.topology == "2x2"
+    mock_tpu_client.create_node.assert_called_once()
+    operation.result.assert_called_once()