forked from luopingyi/xc_agent_platform_demo
62 lines
3.2 KiB
Python
62 lines
3.2 KiB
Python
|
import io
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import unittest
|
||
|
|
import urllib.error
|
||
|
|
from unittest.mock import patch, MagicMock
|
||
|
|
from src.readonly_probe import probe_readonly
|
||
|
|
from scripts.probe_readonly import main
|
||
|
|
|
||
|
|
|
||
|
|
class TestReadonlyProbe(unittest.TestCase):
|
||
|
|
def test_missing_token_no_network(self):
|
||
|
|
with patch.dict(os.environ, {}, clear=True), patch('urllib.request.urlopen') as call:
|
||
|
|
self.assertFalse(probe_readonly()['auth_ok'])
|
||
|
|
with patch('sys.stdout', io.StringIO()):
|
||
|
|
self.assertEqual(main(), 1)
|
||
|
|
call.assert_not_called()
|
||
|
|
|
||
|
|
def test_success_and_exact_request(self):
|
||
|
|
response = MagicMock()
|
||
|
|
response.__enter__.return_value = response
|
||
|
|
response.status = 200
|
||
|
|
response.read.return_value = json.dumps({'code': 0, 'data': {
|
||
|
|
'records': [{'modelId': 'secret-value'}], 'total': 10}}).encode()
|
||
|
|
with patch.dict(os.environ, {'EXTERNAL_SERVICE_TOKEN': 'secret-value'}, clear=True), patch('urllib.request.urlopen', return_value=response) as call:
|
||
|
|
result = probe_readonly()
|
||
|
|
request = call.call_args.args[0]
|
||
|
|
self.assertEqual(request.get_method(), 'GET')
|
||
|
|
self.assertEqual(request.full_url, 'https://modelhub.org.cn/api/adapt/task/page?onlyMine=true&pageSize=5¤t=1')
|
||
|
|
self.assertEqual(request.get_header('Xc-token'), 'secret-value')
|
||
|
|
self.assertTrue(result['auth_ok'])
|
||
|
|
self.assertEqual(result['record_summary'], {'count': 1, 'total': 10})
|
||
|
|
self.assertNotIn('secret-value', json.dumps(result))
|
||
|
|
|
||
|
|
def test_errors_no_fallback_or_echo(self):
|
||
|
|
for code in (40100, 40101, 50000):
|
||
|
|
error = urllib.error.HTTPError('https://example.test', 403, '', {}, io.BytesIO(json.dumps({'code': code, 'message': 'secret-value'}).encode()))
|
||
|
|
with patch.dict(os.environ, {'XC_TOKEN': 'secret-value'}, clear=True), patch('urllib.request.urlopen', side_effect=error) as call:
|
||
|
|
result = probe_readonly()
|
||
|
|
self.assertEqual(call.call_count, 1)
|
||
|
|
self.assertFalse(result['auth_ok'])
|
||
|
|
self.assertEqual(result['platform_code'], code)
|
||
|
|
self.assertIn(str(code) if code != 50000 else 'other', result['error_attribution'])
|
||
|
|
self.assertNotIn('secret-value', json.dumps(result))
|
||
|
|
|
||
|
|
def test_transport_and_invalid_json(self):
|
||
|
|
for error in (OSError('secret-value'), urllib.error.URLError('secret-value')):
|
||
|
|
with patch.dict(os.environ, {'XC_TOKEN': 'secret-value'}, clear=True), patch('urllib.request.urlopen', side_effect=error):
|
||
|
|
self.assertEqual(probe_readonly()['error_attribution'], 'transport/timeout')
|
||
|
|
response = MagicMock()
|
||
|
|
response.__enter__.return_value = response
|
||
|
|
response.status, response.read.return_value = 200, b'invalid secret-value'
|
||
|
|
with patch.dict(os.environ, {'XC_TOKEN': 'secret-value'}, clear=True), patch('urllib.request.urlopen', return_value=response):
|
||
|
|
self.assertFalse(probe_readonly()['auth_ok'])
|
||
|
|
|
||
|
|
def test_import_no_network(self):
|
||
|
|
import importlib
|
||
|
|
import src.readonly_probe
|
||
|
|
with patch('urllib.request.urlopen') as call:
|
||
|
|
importlib.reload(src.readonly_probe)
|
||
|
|
call.assert_not_called()
|