forked from luopingyi/xc_agent_platform_demo
56 lines
3.0 KiB
Python
56 lines
3.0 KiB
Python
import json
|
|
import os
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
from src.live_queue import read_queue
|
|
|
|
|
|
class TestLiveQueue(unittest.TestCase):
|
|
def response(self, body):
|
|
r = MagicMock()
|
|
r.__enter__.return_value = r
|
|
r.status = 200
|
|
r.read.return_value = json.dumps(body).encode()
|
|
return r
|
|
|
|
def test_pagination_and_status(self):
|
|
pages = [self.response({'code': 0, 'data': {'total': '3', 'records': [
|
|
{'status': 'waiting', 'modelId': 'org/a', 'gpuType': 'Biren_166m'}, {'status': 'success', 'verifyResult': -1}]}}),
|
|
self.response({'code': 0, 'data': {'total': '3', 'records': [{'status': '等待中'}]}})]
|
|
with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}, clear=True), patch('urllib.request.urlopen', side_effect=pages) as call:
|
|
count, why, models = read_queue()
|
|
self.assertEqual(count, 2)
|
|
self.assertIn(('org/a', 'Biren_166m'), models)
|
|
self.assertIn('current=2', call.call_args.args[0].full_url)
|
|
|
|
def test_fail_closed(self):
|
|
for body in ({'code': 0, 'data': {'total': 'bad', 'records': []}}, {'code': 40100}, {'code': 40101}, {'code': 0, 'data': {'total': 1, 'records': []}}, []):
|
|
with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}, clear=True), patch('urllib.request.urlopen', return_value=self.response(body)):
|
|
self.assertIsNone(read_queue()[0])
|
|
with patch.dict(os.environ, {}, clear=True), patch('urllib.request.urlopen') as call:
|
|
self.assertIsNone(read_queue()[0])
|
|
call.assert_not_called()
|
|
|
|
def test_page_limit_and_transport(self):
|
|
with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}, clear=True):
|
|
with patch('urllib.request.urlopen', return_value=self.response({'code': 0, 'data': {'total': 21, 'records': [{'status': 'waiting'}]}})) as call:
|
|
self.assertIsNone(read_queue()[0])
|
|
self.assertEqual(call.call_count, 20)
|
|
with patch('urllib.request.urlopen', side_effect=OSError('synthetic-secret')):
|
|
self.assertNotIn('synthetic-secret', str(read_queue()))
|
|
|
|
def test_invalid_json_http_and_auth_header(self):
|
|
with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}, clear=True):
|
|
r = self.response({})
|
|
r.read.return_value = b'invalid synthetic-secret'
|
|
with patch('urllib.request.urlopen', return_value=r):
|
|
self.assertEqual(read_queue()[1], 'malformed_json')
|
|
r.status = 503
|
|
with patch('urllib.request.urlopen', return_value=r):
|
|
self.assertIsNone(read_queue()[0])
|
|
with patch('urllib.request.urlopen', return_value=self.response({'code': 0, 'data': {'total': 0, 'records': []}})) as call:
|
|
self.assertEqual(read_queue()[0], 0)
|
|
req = call.call_args.args[0]
|
|
self.assertEqual(req.get_header('Xc-token'), 'synthetic-secret')
|
|
self.assertEqual(req.method, 'GET')
|