From cf0a447d941233a9e5dc5a6a1efdf93103f97a26 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Germ=C3=A1n=20Grandas?= Date: Mon, 9 Sep 2024 17:53:23 +0000 Subject: [PATCH] Adding exceptions to test --- applications/rag/tests/test_rag.py | 241 +++++++++++++++-------------- 1 file changed, 127 insertions(+), 114 deletions(-) diff --git a/applications/rag/tests/test_rag.py b/applications/rag/tests/test_rag.py index 3f962f653..83d13087d 100644 --- a/applications/rag/tests/test_rag.py +++ b/applications/rag/tests/test_rag.py @@ -3,121 +3,134 @@ import requests def test_prompts(prompt_url): - testcases = [ - { - "prompt": "List the cast of Squid Game", - }, - { - "prompt": "When was Squid Game released?", - }, - { - "prompt": "What is the rating of Squid Game?", - }, - { - "prompt": "List the cast of Avatar: The Last Airbender", - }, - { - "prompt": "When was Avatar: The Last Airbender added on Netflix?", - }, - { - "prompt": "What is the rating of Avatar: The Last Airbender?", - }, - ] - - for testcase in testcases: - prompt = testcase["prompt"] - - print(f"Testing prompt: {prompt}") - data = {"prompt": prompt} - json_payload = json.dumps(data) - - headers = {'Content-Type': 'application/json'} - response = requests.post(prompt_url, data=json_payload, headers=headers) - response.raise_for_status() - - response = response.json() - print(response) - text = response['response'].get('text') - - print(f"Reply: {text}") - - assert response != None, f"Not response found: {response}" - assert text != None, f"Not text" + try: + testcases = [ + { + "prompt": "List the cast of Squid Game", + }, + { + "prompt": "When was Squid Game released?", + }, + { + "prompt": "What is the rating of Squid Game?", + }, + { + "prompt": "List the cast of Avatar: The Last Airbender", + }, + { + "prompt": "When was Avatar: The Last Airbender added on Netflix?", + }, + { + "prompt": "What is the rating of Avatar: The Last Airbender?", + }, + ] + + for testcase in testcases: + prompt = testcase["prompt"] + + print(f"Testing prompt: {prompt}") + data = {"prompt": prompt} + json_payload = json.dumps(data) + + headers = {'Content-Type': 'application/json'} + response = requests.post(prompt_url, data=json_payload, headers=headers) + response.raise_for_status() + + response = response.json() + print(response) + text = response['response'].get('text') + + print(f"Reply: {text}") + + assert response != None, f"Not response found: {response}" + assert text != None, f"Not text" + except Exception as err: + print(err) + raise err def test_prompts_nlp(prompt_url): - testcases = [ - { - "prompt": "List the cast of Squid Game", - "nlpFilterLevel": "0", - }, - { - "prompt": "Which movie has a building that blows up?", - "nlpFilterLevel": "0", - }, - { - "prompt": "Which movie has a building that blows up?", - "nlpFilterLevel": "50", - }, - { - "prompt": "List the cast of Squid Game", - "nlpFilterLevel": "100", - } - ] - - for testcase in testcases: - prompt = testcase["prompt"] - nlpFilterLevel = testcase["nlpFilterLevel"] - - print(f"Testing prompt: {prompt}") - data = {"prompt": prompt, "nlpFilterLevel": nlpFilterLevel} - json_payload = json.dumps(data) - - headers = {'Content-Type': 'application/json'} - response = requests.post(prompt_url, data=json_payload, headers=headers) - response.raise_for_status() - - response = response.json() - - text = response['response']['text'] - - - print(f"Reply: {text}") - - assert response != None, f"Not response found: {response}" - assert text != None, f"Not text" + try: + testcases = [ + { + "prompt": "List the cast of Squid Game", + "nlpFilterLevel": "0", + }, + { + "prompt": "Which movie has a building that blows up?", + "nlpFilterLevel": "0", + }, + { + "prompt": "Which movie has a building that blows up?", + "nlpFilterLevel": "50", + }, + { + "prompt": "List the cast of Squid Game", + "nlpFilterLevel": "100", + } + ] + + for testcase in testcases: + prompt = testcase["prompt"] + nlpFilterLevel = testcase["nlpFilterLevel"] + + print(f"Testing prompt: {prompt}") + data = {"prompt": prompt, "nlpFilterLevel": nlpFilterLevel} + json_payload = json.dumps(data) + + headers = {'Content-Type': 'application/json'} + response = requests.post(prompt_url, data=json_payload, headers=headers) + response.raise_for_status() + + response = response.json() + + text = response['response']['text'] + + + print(f"Reply: {text}") + + assert response != None, f"Not response found: {response}" + assert text != None, f"Not text" + except Exception as err: + print(err) + raise err def test_prompts_dlp(prompt_url): - testcases = [ - { - "prompt": "who worked with Robert De Niro and name one film they collaborated?", - "inspectTemplate": "projects/gke-ai-eco-dev/locations/global/inspectTemplates/DO-NOT-DELETE-e2e-test-inspect-template", - "deidentifyTemplate": "projects/gke-ai-eco-dev/locations/global/deidentifyTemplates/DO-NOT-DELETE-e2e-test-de-identify-template", - }, - ] - - for testcase in testcases: - prompt = testcase["prompt"] - inspectTemplate = testcase["inspectTemplate"] - deidentifyTemplate = testcase["deidentifyTemplate"] - - print(f"Testing prompt: {prompt}") - data = {"prompt": prompt, "inspectTemplate": inspectTemplate, "deidentifyTemplate": deidentifyTemplate} - json_payload = json.dumps(data) - - headers = {'Content-Type': 'application/json'} - response = requests.post(prompt_url, data=json_payload, headers=headers) - response.raise_for_status() - - response = response.json() - text = response['response']['text'] - - - print(f"Reply: {text}") - - assert response != None, f"Not response found: {response}" - assert text != None, f"Not text" - -prompt_url = sys.argv[1] -test_prompts(prompt_url) -test_prompts_nlp(prompt_url) -test_prompts_dlp(prompt_url) \ No newline at end of file + try: + testcases = [ + { + "prompt": "who worked with Robert De Niro and name one film they collaborated?", + "inspectTemplate": "projects/gke-ai-eco-dev/locations/global/inspectTemplates/DO-NOT-DELETE-e2e-test-inspect-template", + "deidentifyTemplate": "projects/gke-ai-eco-dev/locations/global/deidentifyTemplates/DO-NOT-DELETE-e2e-test-de-identify-template", + }, + ] + + for testcase in testcases: + prompt = testcase["prompt"] + inspectTemplate = testcase["inspectTemplate"] + deidentifyTemplate = testcase["deidentifyTemplate"] + + print(f"Testing prompt: {prompt}") + data = {"prompt": prompt, "inspectTemplate": inspectTemplate, "deidentifyTemplate": deidentifyTemplate} + json_payload = json.dumps(data) + + headers = {'Content-Type': 'application/json'} + response = requests.post(prompt_url, data=json_payload, headers=headers) + response.raise_for_status() + + response = response.json() + text = response['response']['text'] + + + print(f"Reply: {text}") + + assert response != None, f"Not response found: {response}" + assert text != None, f"Not text" + except Exception as err: + print(err) + raise err + +if __name__ = "__main__": + prompt_url = sys.argv[1] + test_prompts(prompt_url) + test_prompts_nlp(prompt_url) + test_prompts_dlp(prompt_url) \ No newline at end of file