From ea964e77b8e1c636791c9de026963d0d0730259b Mon Sep 17 00:00:00 2001
From: grydz <grydz@pm.me>
Date: Wed, 11 Dec 2024 16:23:59 +0400
Subject: [PATCH] Fix: tests of Yao's Millionaires example

---
 examples/yaos_millionaires/README.md         |  2 +-
 examples/yaos_millionaires/tests/conftest.py | 26 ++++++--------
 examples/yaos_millionaires/tests/test_app.py | 36 ++++++++------------
 3 files changed, 27 insertions(+), 37 deletions(-)

diff --git a/examples/yaos_millionaires/README.md b/examples/yaos_millionaires/README.md
index db46ff6..2ef7073 100644
--- a/examples/yaos_millionaires/README.md
+++ b/examples/yaos_millionaires/README.md
@@ -95,7 +95,7 @@ $ cenclave seal --input secrets_to_seal.json --receiver-enclave ratls.pem --outp
 ### Run your application
 
 ```console
-$ cenclave run  yaos_millionaires
+$ cenclave run yaos_millionaires
 ```
 
 Note: if the code is encrypted, sealed secrets must be added to the command with `--sealed-secrets sealed_secrets.json.enc`.
diff --git a/examples/yaos_millionaires/tests/conftest.py b/examples/yaos_millionaires/tests/conftest.py
index 01aead6..9fe1b49 100644
--- a/examples/yaos_millionaires/tests/conftest.py
+++ b/examples/yaos_millionaires/tests/conftest.py
@@ -6,27 +6,14 @@
 from typing import Optional
 
 import pytest
+import requests
 from intel_sgx_ra.ratls import get_server_certificate, url_parse
 
 
 @pytest.fixture(scope="module")
 def url() -> str:
     """Get the url of the web app."""
-    return os.getenv("TEST_REMOTE_URL", "http://localhost:5000")
-
-
-@pytest.fixture(scope="module")
-def secret_json() -> Optional[Path]:
-    """Get the secret.json path."""
-    e = os.getenv("TEST_SECRET_JSON")
-    return Path(e) if e else None
-
-
-@pytest.fixture(scope="module")
-def sealed_secret_json() -> Optional[Path]:
-    """Get the sealed_secret.json path."""
-    e = os.getenv("TEST_SEALED_SECRET_JSON")
-    return Path(e) if e else None
+    return os.getenv("TEST_URL", "http://localhost:5000")
 
 
 @pytest.fixture(scope="module")
@@ -57,6 +44,15 @@ def certificate(url, workspace) -> Optional[Path]:
     return cert_path
 
 
+@pytest.fixture(scope="module")
+def session(certificate) -> requests.Session:
+    session = requests.Session()
+    if certificate:
+        session.verify = f"{certificate}"
+
+    return session
+
+
 @pytest.fixture(scope="module")
 def keypair1_path() -> bytes:
     """Path of the participant 1 keypair."""
diff --git a/examples/yaos_millionaires/tests/test_app.py b/examples/yaos_millionaires/tests/test_app.py
index ae679b2..7477e46 100644
--- a/examples/yaos_millionaires/tests/test_app.py
+++ b/examples/yaos_millionaires/tests/test_app.py
@@ -7,19 +7,15 @@
 from cenclave_lib_crypto.seal_box import unseal
 
 
-def test_health(url, certificate):
+def test_health(url, session):
     """Test healthcheck endpoint."""
-    response = requests.get(f"{url}/health", timeout=10, verify=certificate)
+    response = session.get(f"{url}/health", timeout=10)
     assert response.status_code == 200
 
 
-def test_participants(url, certificate, pk1, pk1_b64, pk2, pk2_b64):
+def test_participants(url, session, pk1, pk1_b64, pk2, pk2_b64):
     """Test participants endpoint."""
-    response = requests.get(
-        f"{url}/participants",
-        timeout=10,
-        verify=certificate,
-    )
+    response = session.get(f"{url}/participants", timeout=10)
 
     assert response.status_code == 200
 
@@ -35,32 +31,32 @@ def test_participants(url, certificate, pk1, pk1_b64, pk2, pk2_b64):
     assert pk2_b64.decode("utf-8") in result["participants"]
 
 
-def test_richest(url, certificate, pk1_b64, sk1, pk2_b64, sk2):
+def test_richest(url, session, pk1_b64, sk1, pk2_b64, sk2):
+    # reset first
+    response = session.delete(f"{url}", timeout=10)
+    assert response.status_code == 200
+
     n_b64 = base64.b64encode(struct.pack("<d", float(97))).decode("utf-8")
-    response = requests.post(
+    response = session.post(
         url,
         json={"pk": pk1_b64.decode("utf-8"), "data": {"encrypted": False, "n": n_b64}},
         timeout=10,
-        verify=certificate,
     )
 
     assert response.status_code == 200
 
     n_b64 = base64.b64encode(struct.pack("<d", float(97.1))).decode("utf-8")
-    response = requests.post(
+    response = session.post(
         url,
         json={"pk": pk2_b64.decode("utf-8"), "data": {"encrypted": False, "n": n_b64}},
         timeout=10,
-        verify=certificate,
     )
 
     assert response.status_code == 200
 
     # test result with pk1
-    response = requests.post(
-        url=f"{url}/richest",
-        json={"recipient_pk": pk1_b64.decode("utf-8")},
-        verify=certificate,
+    response = session.post(
+        url=f"{url}/richest", json={"recipient_pk": pk1_b64.decode("utf-8")}
     )
 
     assert response.status_code == 200
@@ -77,10 +73,8 @@ def test_richest(url, certificate, pk1_b64, sk1, pk2_b64, sk2):
     assert base64.b64encode(pk_winner) == pk2_b64
 
     # test result with pk2
-    response = requests.post(
-        url=f"{url}/richest",
-        json={"recipient_pk": pk2_b64.decode("utf-8")},
-        verify=certificate,
+    response = session.post(
+        url=f"{url}/richest", json={"recipient_pk": pk2_b64.decode("utf-8")}
     )
 
     assert response.status_code == 200