hf test updated (docker works now)

tests/test_hf_api.py

@@ -12,10 +12,15 @@ def get_final_answer(text: str) -> int | None:
 BASE_URL = "https://tcmmichaelb139-evolutiontransformer.hf.space"
 
 
-def await_task_completion(task_id, timeout=60):
+@pytest.fixture
+def session():
+    return requests.Session()
+
+
+def await_task_completion(session, task_id, timeout=60):
     start_time = time.time()
     while time.time() - start_time < timeout:
-        status_response = requests.get(f"{BASE_URL}/tasks/{task_id}")
+        status_response = session.get(f"{BASE_URL}/tasks/{task_id}")
 
         print(status_response.json())
 
@@ -36,11 +41,11 @@ def await_task_completion(task_id, timeout=60):
     return None
 
 
-def test_generate_endpoint_svamp():
+def test_generate_endpoint_svamp(session):
     """
     Tests inference on svamp
     """
-    response = requests.post(
+    response = session.post(
         f"{BASE_URL}/generate",
         json={
             "model_name": "svamp",
@@ -56,7 +61,7 @@ def test_generate_endpoint_svamp():
     assert "task_id" in data
     task_id = data["task_id"]
 
-    final_result = await_task_completion(task_id)
+    final_result = await_task_completion(session, task_id)
 
     assert "response" in final_result
     output_text = final_result["response"]
@@ -65,12 +70,12 @@ def test_generate_endpoint_svamp():
     assert answer == 14
 
 
-def test_merge_then_inference_svamp_1():
+def test_merge_then_inference_svamp_1(session):
     """
     Tests merging then inference for svamp dataset
     """
 
-    merge_response = requests.post(
+    merge_response = session.post(
         f"{BASE_URL}/merge",
         json={
             "model1_name": "svamp",
@@ -87,12 +92,12 @@ def test_merge_then_inference_svamp_1():
     assert "task_id" in merge_data
     merge_task_id = merge_data["task_id"]
 
-    merge_status_data = await_task_completion(merge_task_id)
+    merge_status_data = await_task_completion(session, merge_task_id)
     model_name = merge_status_data["response"]
 
     time.sleep(5)
 
-    generate_response = requests.post(
+    generate_response = session.post(
         f"{BASE_URL}/generate",
         json={
             "model_name": model_name,
@@ -107,7 +112,7 @@ def test_merge_then_inference_svamp_1():
     assert "task_id" in generate_data
     generate_task_id = generate_data["task_id"]
 
-    final_result = await_task_completion(generate_task_id)
+    final_result = await_task_completion(session, generate_task_id)
 
     assert "response" in final_result
     output_text = final_result["response"]
@@ -116,12 +121,12 @@ def test_merge_then_inference_svamp_1():
     assert answer == 14
 
 
-def test_merge_then_inference_svamp_2():
+def test_merge_then_inference_svamp_2(session):
     """
     Tests merging then inference for svamp dataset
     """
 
-    merge_repsonse = requests.post(
+    merge_repsonse = session.post(
         f"{BASE_URL}/merge",
         json={
             "model1_name": "svamp",
@@ -138,11 +143,11 @@ def test_merge_then_inference_svamp_2():
     assert "task_id" in merge_data
     merge_task_id = merge_data["task_id"]
 
-    merge_status_data = await_task_completion(merge_task_id)
+    merge_status_data = await_task_completion(session, merge_task_id)
 
     model_name = merge_status_data["response"]
 
-    merge_response2 = requests.post(
+    merge_response2 = session.post(
         f"{BASE_URL}/merge",
         json={
             "model1_name": model_name,
@@ -158,12 +163,12 @@ def test_merge_then_inference_svamp_2():
     merge_data2 = merge_response2.json()
     assert "task_id" in merge_data2
     merge_task_id2 = merge_data2["task_id"]
-    merge_status_data2 = await_task_completion(merge_task_id2)
+    merge_status_data2 = await_task_completion(session, merge_task_id2)
     model_name2 = merge_status_data2["response"]
 
     time.sleep(5)
 
-    generate_response = requests.post(
+    generate_response = session.post(
         f"{BASE_URL}/generate",
         json={
             "model_name": model_name2,
@@ -178,7 +183,7 @@ def test_merge_then_inference_svamp_2():
     assert "task_id" in generate_data
     generate_task_id = generate_data["task_id"]
 
-    final_result = await_task_completion(generate_task_id)
+    final_result = await_task_completion(session, generate_task_id)
 
     assert "response" in final_result
     output_text = final_result["response"]
@@ -187,12 +192,12 @@ def test_merge_then_inference_svamp_2():
     assert answer == 14
 
 
-def test_merge_two_children_then_merge():
+def test_merge_two_children_then_merge(session):
     """
     Tests creating two children and merging them
     """
 
-    merge_response1 = requests.post(
+    merge_response1 = session.post(
         f"{BASE_URL}/merge",
         json={
             "model1_name": "svamp",
@@ -209,10 +214,10 @@ def test_merge_two_children_then_merge():
     merge_data1 = merge_response1.json()
     assert "task_id" in merge_data1
     merge_task_id1 = merge_data1["task_id"]
-    merge_status_data1 = await_task_completion(merge_task_id1)
+    merge_status_data1 = await_task_completion(session, merge_task_id1)
     child1_name = merge_status_data1["response"]
 
-    merge_response2 = requests.post(
+    merge_response2 = session.post(
         f"{BASE_URL}/merge",
         json={
             "model1_name": "svamp",
@@ -229,10 +234,10 @@ def test_merge_two_children_then_merge():
     merge_data2 = merge_response2.json()
     assert "task_id" in merge_data2
     merge_task_id2 = merge_data2["task_id"]
-    merge_status_data2 = await_task_completion(merge_task_id2)
+    merge_status_data2 = await_task_completion(session, merge_task_id2)
     child2_name = merge_status_data2["response"]
 
-    merge_response3 = requests.post(
+    merge_response3 = session.post(
         f"{BASE_URL}/merge",
         json={
             "model1_name": child1_name,
@@ -248,12 +253,12 @@ def test_merge_two_children_then_merge():
     merge_data3 = merge_response3.json()
     assert "task_id" in merge_data3
     merge_task_id3 = merge_data3["task_id"]
-    merge_status_data3 = await_task_completion(merge_task_id3)
+    merge_status_data3 = await_task_completion(session, merge_task_id3)
     final_model_name = merge_status_data3["response"]
 
     time.sleep(5)
 
-    generate_response = requests.post(
+    generate_response = session.post(
         f"{BASE_URL}/generate",
         json={
             "model_name": final_model_name,
@@ -268,7 +273,7 @@ def test_merge_two_children_then_merge():
     assert "task_id" in generate_data
     generate_task_id = generate_data["task_id"]
 
-    final_result = await_task_completion(generate_task_id)
+    final_result = await_task_completion(session, generate_task_id)
 
     assert "response" in final_result
     output_text = final_result["response"]
@@ -277,12 +282,12 @@ def test_merge_two_children_then_merge():
     assert answer == 14
 
 
-def test_merge_fail():
+def test_merge_fail(session):
     """
     Tests merging with too many layers
     """
 
-    merge_repsonse = requests.post(
+    merge_repsonse = session.post(
         f"{BASE_URL}/merge",
         json={
             "model1_name": "svamp",
@@ -299,7 +304,7 @@ def test_merge_fail():
     assert "task_id" in merge_data
     merge_task_id = merge_data["task_id"]
 
-    merge_status_data = await_task_completion(merge_task_id)
+    merge_status_data = await_task_completion(session, merge_task_id)
     assert "response" not in merge_status_data
     assert "error" in merge_status_data
     assert "Layer recipe too long" in merge_status_data["error"]