risingwave_license/
rwu.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// Tests below only work in debug mode.
16#[cfg(debug_assertions)]
17#[cfg(test)]
18mod tests {
19    const GB: u64 = 1024 * 1024 * 1024;
20
21    use expect_test::expect;
22    use risingwave_pb::common::ClusterResource;
23    use thiserror_ext::AsReport as _;
24
25    use crate::{Feature, LicenseKey, LicenseManager, TEST_ALL_LICENSE_KEY_CONTENT};
26
27    fn do_test(key: &str, cpu_core_count: u64, memory_bytes: u64, expect: expect_test::Expect) {
28        let manager = LicenseManager::new();
29        manager.refresh(LicenseKey(key));
30        manager.update_cluster_resource(ClusterResource {
31            total_memory_bytes: memory_bytes,
32            total_cpu_cores: cpu_core_count,
33        });
34
35        match Feature::TestDummy.check_available_with(&manager) {
36            Ok(_) => expect.assert_eq("ok"),
37            Err(error) => expect.assert_eq(&error.to_report_string()),
38        }
39    }
40
41    #[test]
42    fn test_no_limit() {
43        do_test(TEST_ALL_LICENSE_KEY_CONTENT, 114514, 0, expect!["ok"]);
44        do_test(TEST_ALL_LICENSE_KEY_CONTENT, 0, 114514 * GB, expect!["ok"]);
45    }
46
47    #[test]
48    fn test_invalid_license_key() {
49        const KEY: &str = "invalid";
50
51        do_test(
52            KEY,
53            0,
54            0,
55            expect![
56                "feature TestDummy is not available due to license error: invalid license key: InvalidToken"
57            ],
58        );
59        do_test(
60            KEY,
61            114514,
62            0,
63            expect![
64                "feature TestDummy is not available due to license error: invalid license key: InvalidToken"
65            ],
66        );
67    }
68
69    const KEY_32: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
70          eyJzdWIiOiJwYWlkLXRlc3QtMzIiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwidGllciI6InBhaWQiLCJleHAiOjIxNTA0OTU5OTksImlhdCI6MTczNzYxMjQ5NSwiY3B1X2NvcmVfbGltaXQiOjMyfQ.\
71          SQpX2Dmon5Mb04VUbHyxsU7owJhcdLZHqUefxAXBwG5AqgKdpfS0XUePW5E4D-EfxtH_cWJiD4QDFsfdRUz88g_n_KvfNUObMW7NV5TUoRs_ImtS4ySugExNX3JzJi71QqgI8kugStQ7uOR9kZ_C-cCc_IG2CwwEmhhW1Ij0vX7qjhG5JNMit_bhxPY7Rh27ppgPTqWxJFTTsw-9B7O5WR_yIlaDjxVzk0ALm_j6DPB249gG3dkeK0rP0AK_ip2cK6iQdy8Cge7ATD6yUh4c_aR6GILDF6-vyB7QdWU6DdQS4KhdkPNWoe_Z9psotcXQJ7NhQ39hk8tdLzmTfGDDBA";
72
73    #[test]
74    fn test_cpu_limit() {
75        do_test(KEY_32, 31, 0, expect!["ok"]);
76        do_test(KEY_32, 32, 0, expect!["ok"]);
77        do_test(
78            KEY_32,
79            33,
80            0,
81            expect![
82                "feature TestDummy is not available due to license error: a valid license key is set, but it is currently not effective because the CPU core in the cluster (33) exceeds the maximum allowed by the license key (32); consider removing some nodes or acquiring a new license key with a higher limit"
83            ],
84        );
85    }
86
87    #[test]
88    fn test_memory_limit() {
89        do_test(KEY_32, 0, 31 * 4 * GB, expect!["ok"]);
90        do_test(KEY_32, 0, 32 * 4 * GB, expect!["ok"]);
91        do_test(
92            KEY_32,
93            0,
94            33 * 4 * GB,
95            expect![
96                "feature TestDummy is not available due to license error: a valid license key is set, but it is currently not effective because the memory in the cluster (132 GiB) exceeds the maximum allowed by the license key (132.00 GiB); consider removing some nodes or acquiring a new license key with a higher limit"
97            ],
98        );
99    }
100}