import base64
from pathlib import Path
import requests
import unittest

# 测试test环境
test_url = 'http://aihub-test.digitalyili.com/aiSquare/openApi/reasoning-services/rlocrxm'
test_header = {
    'Authorization': 'Bearer 9679c2b3-b90b-4029-a3c7-f347b4d242f7'
}

def send_request(image_path,suffix):
    with open(image_path, 'rb') as f:
        img_str: str = base64.encodebytes(f.read()).decode('utf-8')
        r = requests.post(f'{test_url}+{suffix}', json={'image': img_str, 'image_type': '0'}, headers=test_header)
        return r.json()

class TestOcr(unittest.TestCase):

    def _helper(self, image_path, suffix):
        root = Path(__file__).parent
        image_path = str(root / image_path)
        r = send_request(image_path, suffix)
        self.assertIn('result', r, f'{suffix} test error')

    def test_bank(self):
        image_path = './image/bankcard.jpg'
        suffix = '/yhksb/bankcard'
        self._helper(image_path, suffix)

    def test_idcard(self):
        image_path = './image/idcard.jpg'
        suffix = '/sfzsb/idcard'
        self._helper(image_path, suffix)

    def test_regbook(self):
        image_path = './image/regbook.jpg'
        suffix = '/hkbsb/regbook'
        self._helper(image_path, suffix)

    def test_schoolcert(self):
        image_path = './image/schoolcert.jpg'
        suffix = '/schoolcert'
        self._helper(image_path, suffix)





# 测试sb环境
url = 'http://aihub.digitalyili.com/aiSquare/openApi/reasoning-services/rlocrxm'
header = {
    'Content-Type': 'application/json',
    'Authorization': 'Bearer dcae8cc6-0e49-4db8-a2d2-94ef84da3636'
}

def send_sb_request(image_path,suffix):
    with open(image_path, 'rb') as f:
        img_str: str = base64.encodebytes(f.read()).decode('utf-8')
        r = requests.post(f'{url}{suffix}', json={'image': img_str, 'image_type': '0'}, headers=header)
        return r.json()

class TestSbOcr(unittest.TestCase):

    def _helper1(self, image_path, suffix):
        root = Path(__file__).parent
        image_path = str(root / image_path)
        r = send_sb_request(image_path, suffix)
        self.assertIn('result', r, f'{suffix} test error')

    def test_bank(self):
        image_path = './image/bankcard.jpg'
        suffix = '/yhksb/bankcard'
        self._helper1(image_path, suffix)

    def test_idcard(self):
        image_path = './image/idcard.jpg'
        suffix = '/sfzsb/idcard'
        self._helper1(image_path, suffix)

    def test_regbook(self):
        image_path = './image/regbook.jpg'
        suffix = '/hkbsb/regbook'
        self._helper1(image_path, suffix)

    def test_schoolcert(self):
        image_path = './image/schoolcert.jpg'
        suffix = '/schoolcert'
        self._helper1(image_path, suffix)