wavespeed-python/tests/test_config.py at e6ff395d878d6f3690a35f58fccf690116c0b7a8 · WaveSpeedAI/wavespeed-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
"""Tests for the config module."""

import unittest

from wavespeed.config import _resolve_runpod_url, _resolve_waverless_url, serverless


class TestResolveRunpodUrl(unittest.TestCase):
    """Tests for the _resolve_runpod_url function."""

    def test_replaces_runpod_pod_id(self):
        """Test that $RUNPOD_POD_ID is replaced with pod_id."""
        template = "https://api.runpod.ai/v2/endpoint/job-done/$RUNPOD_POD_ID"
        result = _resolve_runpod_url(template, "my-pod-123")
        self.assertEqual(
            result, "https://api.runpod.ai/v2/endpoint/job-done/my-pod-123"
        )

    def test_preserves_id_placeholder(self):
        """Test that $ID is NOT replaced - it's for job ID at runtime."""
        template = "https://api.runpod.ai/v2/endpoint/job-done/$RUNPOD_POD_ID/$ID"
        result = _resolve_runpod_url(template, "my-pod-123")
        self.assertEqual(
            result, "https://api.runpod.ai/v2/endpoint/job-done/my-pod-123/$ID"
        )

    def test_handles_none_template(self):
        """Test that None template returns None."""
        result = _resolve_runpod_url(None, "my-pod-123")
        self.assertIsNone(result)

    def test_no_placeholders(self):
        """Test URL without any placeholders."""
        template = "https://api.example.com/endpoint"
        result = _resolve_runpod_url(template, "my-pod-123")
        self.assertEqual(result, "https://api.example.com/endpoint")


class TestResolveWaverlessUrl(unittest.TestCase):
    """Tests for the _resolve_waverless_url function."""

    def test_replaces_waverless_pod_id_placeholder(self):
        """Test that $WAVERLESS_POD_ID is replaced with pod_id."""
        template = "https://api.wavespeed.ai/v2/test/job-take/$WAVERLESS_POD_ID"
        result = _resolve_waverless_url(template, "my-pod-123")
        self.assertEqual(
            result, "https://api.wavespeed.ai/v2/test/job-take/my-pod-123"
        )

    def test_preserves_id_placeholder(self):
        """Test that $ID is NOT replaced - it's for job/worker ID at runtime."""
        template = "https://api.wavespeed.ai/v2/test/job-done/$WAVERLESS_POD_ID/$ID"
        result = _resolve_waverless_url(template, "my-pod-123")
        self.assertEqual(
            result, "https://api.wavespeed.ai/v2/test/job-done/my-pod-123/$ID"
        )

    def test_handles_none_template(self):
        """Test that None template returns None."""
        result = _resolve_waverless_url(None, "my-pod-123")
        self.assertIsNone(result)

    def test_no_placeholders(self):
        """Test URL without any placeholders."""
        template = "https://api.example.com/endpoint"
        result = _resolve_waverless_url(template, "my-pod-123")
        self.assertEqual(result, "https://api.example.com/endpoint")


class TestServerlessConfig(unittest.TestCase):
    """Tests for serverless config loading."""

    def test_serverless_has_expected_attributes(self):
        """Test that serverless config has all expected attributes."""
        self.assertTrue(hasattr(serverless, "pod_id"))
        self.assertTrue(hasattr(serverless, "api_key"))
        self.assertTrue(hasattr(serverless, "job_get_url"))
        self.assertTrue(hasattr(serverless, "job_done_url"))
        self.assertTrue(hasattr(serverless, "job_stream_url"))
        self.assertTrue(hasattr(serverless, "ping_url"))


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