runpod-python/tests/test_api/test_mutation_endpoints.py at 9fc995b5e1c03be5a17c58e5e20f843ac7f220cb · 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
"""Tests for the endpoint mutation generation."""

import unittest

from runpod.api.mutations.endpoints import generate_endpoint_mutation


class TestGenerateEndpointMutation(unittest.TestCase):
    """Tests for the endpoint mutation generation."""

    def test_required_fields(self):
        """Test the required fields."""
        result = generate_endpoint_mutation("test_name", "test_template_id")
        self.assertIn('name: "test_name"', result)
        self.assertIn('templateId: "test_template_id"', result)
        self.assertIn('gpuIds: "AMPERE_16"', result)  # Default value
        self.assertIn('networkVolumeId: ""', result)  # Default value
        self.assertIn('locations: ""', result)  # Default value

    def test_all_fields(self):
        """Test all the fields."""
        result = generate_endpoint_mutation(
            "test_name",
            "test_template_id",
            "AMPERE_20",
            "test_volume_id",
            "US_WEST",
            10,
            "WORKER_COUNT",
            5,
            2,
            4,
            True,
        )
        self.assertIn('name: "test_name"', result)
        self.assertIn('templateId: "test_template_id"', result)
        self.assertIn('gpuIds: "AMPERE_20"', result)
        self.assertIn('networkVolumeId: "test_volume_id"', result)
        self.assertIn('locations: "US_WEST"', result)
        self.assertIn("idleTimeout: 10", result)
        self.assertIn('scalerType: "WORKER_COUNT"', result)
        self.assertIn("scalerValue: 5", result)
        self.assertIn("workersMin: 2", result)
        self.assertIn("workersMax: 4", result)
        self.assertIn("flashBootType: FLASHBOOT", result)