test_embedding.py

import json
import time
from websocket import create_connection
#NOTE: needs installation of pytest and pytest-benchamark libraries

sentences = ['How do you treat a migraine?']

uri = "ws://{}:{}{}".format('node1.naint.tech', 8383, '/services/bio_clinical_bert/')
ws = create_connection(uri)

def get_embedding(sent=sentences[0]):
    data = {"sentences": sentences, "event": "bio_clinical_bert", "clean_embedding": True}
    ws.send(json.dumps(data))
    responce = json.loads(ws.recv())
    return responce


def test_embedding(benchmark):
    benchmark(get_embedding)