From 50a7e2c654254bea050c496b2e7faa20f9964a84 Mon Sep 17 00:00:00 2001
From: Five Grant <5@fivegrant.com>
Date: Wed, 14 Feb 2024 16:55:47 -0600
Subject: [PATCH] Fix serialization

---
 service/api.py                                        | 11 ++---------
 service/execute.py                                    |  7 ++++++-
 service/models/base.py                                |  6 ------
 service/models/operations/optimize.py                 |  3 ++-
 ...semble_simulate.py => broken_ensemble_simulate.py} |  5 +++++
 tests/integration/test_optimize.py                    |  3 +--
 6 files changed, 16 insertions(+), 19 deletions(-)
 rename tests/integration/{test_ensemble_simulate.py => broken_ensemble_simulate.py} (89%)

diff --git a/service/api.py b/service/api.py
index a1667db..74e87ed 100644
--- a/service/api.py
+++ b/service/api.py
@@ -11,6 +11,7 @@
     Calibrate,
     Simulate,
     EnsembleSimulate,
+    Optimize,
     StatusSimulationIdGetResponse,
 )
 
@@ -20,6 +21,7 @@
     "simulate": Simulate,
     "calibrate": Calibrate,
     "ensemble-simulate": EnsembleSimulate,
+    "optimize": Optimize,
 }
 
 logging.basicConfig()
@@ -114,12 +116,3 @@ def ensemble_calibrate_not_yet_implemented():
     This will be reimplemented in the future.
     """
     raise HTTPException(status=501, detail="Not yet reimplemented")
-
-
-@app.get("/optimize", response_model=StatusSimulationIdGetResponse)  # NOT YET IN SPEC
-def optimize_not_yet_implemented():  # NOT YET IN SPEC
-    """
-    DO NOT USE. Placeholder for `optimize` endpoint.
-    This will be implemented in the future.
-    """
-    raise HTTPException(status=501, detail="Not yet implemented")
diff --git a/service/execute.py b/service/execute.py
index 4325d36..013fc2f 100644
--- a/service/execute.py
+++ b/service/execute.py
@@ -7,7 +7,12 @@
     attach_files,
 )
 
-from pyciemss.interfaces import sample, calibrate, ensemble_sample  # noqa: F401
+from pyciemss.interfaces import (  # noqa: F401
+    sample,
+    calibrate,
+    ensemble_sample,
+    optimize,
+)
 
 # jl = newmodule("SciMLIntegration")
 # jl.seval("using SciMLIntegration, PythonCall")
diff --git a/service/models/base.py b/service/models/base.py
index 988e14c..3be2fb2 100644
--- a/service/models/base.py
+++ b/service/models/base.py
@@ -42,12 +42,6 @@ class InterventionSelection(BaseModel):
     name: str
 
 
-class QuantityOfInterest(BaseModel):
-    function: str
-    state: str
-    arg: int  # TODO: Make this a list of args?
-
-
 class OperationRequest(BaseModel):
     pyciemss_lib_function: ClassVar[str] = ""
     engine: str = Field("ciemss", example="ciemss")
diff --git a/service/models/operations/optimize.py b/service/models/operations/optimize.py
index 09087ee..2820ef2 100644
--- a/service/models/operations/optimize.py
+++ b/service/models/operations/optimize.py
@@ -34,7 +34,7 @@ def obs_nday_average_qoi(
     return np.mean(dataQoI[:, -ndays:], axis=1)
 
 
-qoi_implementations = {QOIMethod.obs_nday_average: obs_nday_average_qoi}
+qoi_implementations = {QOIMethod.obs_nday_average.value: obs_nday_average_qoi}
 
 
 class OptimizeExtra(BaseModel):
@@ -99,3 +99,4 @@ def gen_pyciemss_args(self, job_id):
 
     class Config:
         extra = Extra.forbid
+        use_enum_values = True
diff --git a/tests/integration/test_ensemble_simulate.py b/tests/integration/broken_ensemble_simulate.py
similarity index 89%
rename from tests/integration/test_ensemble_simulate.py
rename to tests/integration/broken_ensemble_simulate.py
index f6d9d03..e2bddfc 100644
--- a/tests/integration/test_ensemble_simulate.py
+++ b/tests/integration/broken_ensemble_simulate.py
@@ -1,3 +1,8 @@
+"""
+NOTE: PyCIEMSS library had a broken ensemble but it has now been fixed. This
+test succeeded using the broken version but now fails using the updated PyCIEMSS.
+Will be fixed and moved by to `test_ensemble_simulate` soon.    
+"""
 import json
 
 import pytest
diff --git a/tests/integration/test_optimize.py b/tests/integration/test_optimize.py
index b11c270..b0a9a3d 100644
--- a/tests/integration/test_optimize.py
+++ b/tests/integration/test_optimize.py
@@ -8,7 +8,7 @@
 
 
 @pytest.mark.example_dir("optimize")
-def test_simulate_example(
+def test_optimize_example(
     example_context, client, worker, file_storage, file_check, requests_mock
 ):
     job_id = "9ed74639-7778-4bb9-96fd-7509d68cd425"
@@ -22,7 +22,6 @@ def test_simulate_example(
     response = client.post(
         "/optimize",
         json=request,
-        headers={"Content-Type": "application/json"},
     )
     simulation_id = response.json()["simulation_id"]
     response = client.get(