sbt-idp/tests/utils.py

import requests
import traceback
import logging
import time
from requests_toolbelt import  MultipartEncoderMonitor


def login(host, username, password):
    print(host, username, password)
    response = requests.post(f'{host}/api/ctel/login/', json={
        'username': username,
        'password': password
    })
    try:
        token = response.json().get('token', '')
    except:
        logging.error("Failed to login")
        traceback.print_exc()
        raise ValueError("Could not login. Please check host, username or password.")
    if not token:
        logging.error(response.content)
        raise ValueError("Could not get a valid token. Please check host, username or password.")

    return token


def make_sbt_request(host, token, invoice_files=None, imei_files=None, ensure_success=True, sub=None):
    files =  []
    if invoice_files:
        for file in invoice_files:
            files.append(('invoice_file', (file, open(file, "rb"), 'application/octet-stream')))
    if imei_files:
        for file in imei_files:
            files.append(('imei_files', (file, open(file, "rb"), 'application/octet-stream')))
    num_files = len(files)
    files.append(('processType', '12'))
    files.append(('is_test_request', 'True'))
    if sub:
        files.append(("subsidiary", sub))
    start_time = time.time()
    end_of_upload_time = 0

    def upload_callback(monitor):
        nonlocal end_of_upload_time
        if monitor.bytes_read == monitor.len:
            end_of_upload_time = time.time()
    m = MultipartEncoderMonitor.from_fields(
        fields=files,
        callback=upload_callback
    )
    try:
        response = requests.post(f'{host}/api/ctel/images/process_sync/', headers={
            'Authorization': token,
            'Content-Type': m.content_type
        }, data=m, timeout=300)
    except requests.exceptions.Timeout as e:
        print(e)
        traceback.print_exc()
        raise TimeoutError("Timeout when sending request to server.")
    except Exception as e:
        print(e)
        traceback.print_exc()
        raise e
    upload_time = end_of_upload_time - start_time
    process_time = time.time() - start_time - upload_time

    parsed_data = None
    try:
        parsed_data = response.json()
    except:
        pass
    if ensure_success:
        assert upload_time > 0
        assert process_time > 0
        assert num_files > 0
        assert response.status_code == 200

        assert parsed_data is not None
        assert "id" in parsed_data
        assert "pages" in parsed_data
        assert "created_at" in parsed_data
        assert "data" in parsed_data
        assert "files" in parsed_data

    return {
        "upload_time": upload_time,
        "process_time": process_time,
        "data": parsed_data if parsed_data is not None else response.content,
        "num_files": num_files,
    }


def get_field(document, label):
    for field in document["content"]:
        if field["label"] == label:
            return field["value"]
    return None