runpod-python/tests/test_api/test_mutations_pods.py at a3c3a265b53091e6e256f906ca147a2d0c21e575 · runpod/runpod-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
""" Test API Wrapper Pod Mutations """

import unittest

from runpod.api.mutations import pods


class TestPodMutations(unittest.TestCase):
    """Test API Wrapper Pod Mutations"""

    def test_generate_pod_deployment_mutation(self):
        """
        Test generate_pod_deployment_mutation for both GPU and CPU pods
        """
        # Test GPU pod deployment
        gpu_result = pods.generate_pod_deployment_mutation(
            name="test",
            image_name="test_image",
            gpu_type_id="1",
            cloud_type="cloud",
            data_center_id="1",
            country_code="US",
            gpu_count=1,
            volume_in_gb=100,
            container_disk_in_gb=10,
            min_vcpu_count=1,
            min_memory_in_gb=1,
            docker_args="args",
            ports="8080",
            volume_mount_path="/path",
            env={"ENV": "test"},
            support_public_ip=True,
            template_id="abcde",
            allowed_cuda_versions=["11.8", "12.0"],
        )

        # Test CPU pod deployment
        cpu_result = pods.generate_pod_deployment_mutation(
            name="test-cpu",
            image_name="test_image",
            cloud_type="cloud",
            data_center_id="1",
            country_code="US",
            volume_in_gb=100,
            container_disk_in_gb=10,
            min_vcpu_count=2,
            min_memory_in_gb=4,
            docker_args="args",
            ports="8080",
            volume_mount_path="/path",
            env={"ENV": "test"},
            instance_id="cpu3c-2-4"
        )

        # Check GPU pod mutation structure
        self.assertIn("mutation", gpu_result)
        self.assertIn("podFindAndDeployOnDemand", gpu_result)

        # Check CPU pod mutation structure
        self.assertIn("mutation", cpu_result)
        self.assertIn("deployCpuPod", cpu_result)

    def test_generate_pod_deployment_mutation_with_encrypt_volume(self):
        """
        Test generate_pod_deployment_mutation with encrypt_volume parameter
        """
        # Test GPU pod deployment with encryption enabled
        gpu_result_encrypted = pods.generate_pod_deployment_mutation(
            name="test-encrypted",
            image_name="test_image",
            gpu_type_id="1",
            cloud_type="cloud",
            encrypt_volume=True
        )

        # Test CPU pod deployment with encryption enabled
        cpu_result_encrypted = pods.generate_pod_deployment_mutation(
            name="test-cpu-encrypted",
            image_name="test_image",
            cloud_type="cloud",
            instance_id="cpu3c-2-4",
            encrypt_volume=True
        )

        # Test GPU pod deployment with encryption explicitly disabled
        gpu_result_no_encryption = pods.generate_pod_deployment_mutation(
            name="test-no-encryption",
            image_name="test_image",
            gpu_type_id="1",
            cloud_type="cloud",
            encrypt_volume=False
        )

        # Check that encrypted mutations contain encryptVolume: true
        self.assertIn("encryptVolume: true", gpu_result_encrypted)
        self.assertIn("encryptVolume: true", cpu_result_encrypted)

        # Check that non-encrypted mutations do not contain encryptVolume
        self.assertNotIn("encryptVolume", gpu_result_no_encryption)

        # Check basic mutation structure is preserved
        self.assertIn("mutation", gpu_result_encrypted)
        self.assertIn("podFindAndDeployOnDemand", gpu_result_encrypted)
        self.assertIn("mutation", cpu_result_encrypted)
        self.assertIn("deployCpuPod", cpu_result_encrypted)

    def test_generate_pod_stop_mutation(self):
        """
        Test generate_pod_stop_mutation
        """
        result = pods.generate_pod_stop_mutation("pod_id")
        # Here you should check the correct structure of the result
        self.assertIn("mutation", result)

    def test_generate_pod_resume_mutation(self):
        """
        Test generate_pod_resume_mutation
        """
        result = pods.generate_pod_resume_mutation("pod_id", 1)
        # Here you should check the correct structure of the result
        self.assertIn("mutation", result)

    def test_generate_pod_terminate_mutation(self):
        """
        Test generate_pod_terminate_mutation
        """
        result = pods.generate_pod_terminate_mutation("pod_id")
        # Here you should check the correct structure of the result
        self.assertIn("mutation", result)


if __name__ == "__main__":
    unittest.main()