get_subscribe/tests/test_health.py
2026-09-28 02:15:47 +00:00

61 lines
3.7 KiB
Python

import base64
import json
from pathlib import Path
import socket
import tempfile
import unittest
from unittest.mock import Mock, patch
import yaml
import check_nodes
class HealthTests(unittest.TestCase):
def test_endpoints_and_udp_are_not_misclassified(self):
vmess = base64.b64encode(json.dumps({'add': 'example.com', 'port': '443'}).encode()).decode()
ss = base64.b64encode(b'aes-256-gcm:password@example.com:8443').decode()
ssr = base64.urlsafe_b64encode(b'example.com:443:origin:aes-256-cfb:plain:cGFzcw/?x=y').decode()
for node, expected in [({'type': 'ss', 'server': 'example.com', 'port': 443}, ('example.com',443)),
('vmess://' + vmess, ('example.com',443)), ('ss://' + ss, ('example.com',8443)),
('ssr://' + ssr, ('example.com',443)), ('vless://id@[2606:4700:4700::1111]:443', ('2606:4700:4700::1111',443)),
('hysteria2://password@example.com:443',None),
({'type':'tuic', 'server':'example.com', 'port':443},None),
('vless://id@example.com:443?type=quic',None), ('vmess://invalid',None)]:
with self.subTest(node=node):
self.assertEqual(check_nodes.endpoint(node), expected)
def test_private_addresses_are_not_connected(self):
with patch.object(socket, 'getaddrinfo', return_value=[(socket.AF_INET,socket.SOCK_STREAM,6,'',('127.0.0.1',443))]), patch.object(socket, 'socket') as connect:
self.assertEqual(check_nodes.tcp_status(('local.example',443)), 'untested')
connect.assert_not_called()
with patch.object(socket, 'getaddrinfo', side_effect=socket.gaierror):
self.assertEqual(check_nodes.tcp_status(('bad.example',443)), 'unreachable')
def test_tcp_success_and_timeout(self):
resolved = [(socket.AF_INET,socket.SOCK_STREAM,6,'',('1.1.1.1',443))]
with patch.object(socket, 'getaddrinfo', return_value=resolved), patch.object(socket, 'socket') as factory:
self.assertEqual(check_nodes.tcp_status(('public.example',443)), 'reachable')
factory.return_value.__enter__.return_value.connect.assert_called_once_with(('1.1.1.1',443))
factory.return_value.__enter__.return_value.connect.side_effect = socket.timeout
self.assertEqual(check_nodes.tcp_status(('public.example',443)), 'unreachable')
def test_classification_is_complete_and_keeps_original(self):
config = {'proxies': [
{'name':'up','type':'ss','server':'up.example','port':443},
{'name':'down','type':'ss','server':'down.example','port':443},
{'name':'udp','type':'tuic','server':'udp.example','port':443}],
'proxy-groups':[{'name':'select','type':'select','proxies':['up','down','udp']}]}
uris = 'ss://YWVzOnBhc3M=@up.example:443\ntuic://id@udp.example:443\n'
with tempfile.TemporaryDirectory() as directory:
path=Path(directory);source=yaml.safe_dump(config)
(path/'clash.yml').write_text(source);(path/'v2ray.txt').write_text(uris)
with patch.object(check_nodes, 'tcp_status', side_effect=lambda a: 'reachable' if a[0]=='up.example' else 'unreachable') as probe:
result=check_nodes.classify(path)
self.assertEqual(probe.call_count,2)
self.assertEqual(result['counts']['untested'],{'clash':1,'v2ray':1})
self.assertEqual((path/'clash.yml').read_text(),source)
self.assertEqual((path/'v2ray.txt').read_text(),uris)
for status,name in [('reachable','up'),('unreachable','down'),('untested','udp')]:
output=yaml.safe_load((path/status/'clash.yml').read_text())
self.assertEqual(output['proxy-groups'][0]['proxies'],[name])