From ebd8897007c649a876accc3af10949a2d8924a79 Mon Sep 17 00:00:00 2001
From: Brandon Rose <rose.brandon.m@gmail.com>
Date: Wed, 26 Jan 2022 11:35:08 -0600
Subject: [PATCH] stubbed out model testing

---
 api/src/models.py           | 13 +++++++++++-
 api/src/runs.py             | 11 ++++++++++
 api/src/utils.py            | 42 +++++++++++++++++++++++++++++++++++++
 api/validation/RunSchema.py |  7 ++++++-
 4 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/api/src/models.py b/api/src/models.py
index acdf2ba..72d1418 100644
--- a/api/src/models.py
+++ b/api/src/models.py
@@ -16,7 +16,7 @@
 
 from src.settings import settings
 from src.dojo import search_and_scroll, copy_configs, copy_outputfiles, copy_directive, copy_accessory_files
-from src.utils import plugin_action
+from src.utils import plugin_action, run_model_with_defaults
 
 
 router = APIRouter()
@@ -296,3 +296,14 @@ def publish_model(model_id: str, publish_data: ModelSchema.PublishSchema):
         status_code=status.HTTP_200_OK,
         content="Model published",
     )
+
+@router.post("/models/{model_id}/test")
+def test_model(model_id: str):
+    """
+    This endpoint tests a model's functionality within Dojo.
+    """
+    run_id = run_model_with_defaults(model_id)
+    return Response(
+        status_code=status.HTTP_200_OK,
+        content=f"Model test run submitted with run id {run_id}",
+    )
\ No newline at end of file
diff --git a/api/src/runs.py b/api/src/runs.py
index 3e9806c..66c621d 100644
--- a/api/src/runs.py
+++ b/api/src/runs.py
@@ -337,3 +337,14 @@ def update_run(payload: RunSchema.ModelRunSchema):
         headers={"location": f"/api/runs/{run_id}"},
         content=f"Updated run with id = {run_id}",
     )
+
+
+@router.get("/runs/{run_id}/test")
+def test_run_status(run_id: str) -> RunSchema.RunStatusSchema:
+    run = get_run(run_id)
+    status = run.get("attributes",{}).get("status",None)
+    if status:
+        es.index(index="tests", body=body, id=model_id)
+        return status
+    else:
+        return "running"
\ No newline at end of file
diff --git a/api/src/utils.py b/api/src/utils.py
index e5a32a9..42e13f6 100644
--- a/api/src/utils.py
+++ b/api/src/utils.py
@@ -1,5 +1,11 @@
+from dojo.api.src.models import current_milli_time
 from validation import ModelSchema
 from fastapi.logger import logger
+import time
+from elasticsearch import Elasticsearch
+
+from src.settings import settings
+es = Elasticsearch([settings.ELASTICSEARCH_URL], port=settings.ELASTICSEARCH_PORT)
 
 
 def try_parse_int(s: str, default: int = 0) -> int:
@@ -38,6 +44,42 @@ def delete_matching_records_from_model(model_id, record_key, record_test):
     return record_count
 
 
+def run_model_with_defaults(model_id):
+    """
+    This function takes in a model and submits a default run to test that model's functionality
+    """
+
+    from src.models import get_model 
+    from src.runs import create_run, current_milli_time
+    from validation import RunSchema
+
+    model = get_model(model_id)
+
+    params = []
+    for param in model.get("parameters",[]):
+        param_obj = {}
+        param_obj['name'] = param['name']
+        param_obj['value'] = param['default']
+        params.append(param_obj)
+    
+    run_id = f"{model['name']}-{current_milli_time()}"
+
+    run = RunSchema(id=run_id,
+                    model_id=model_id,
+                    model_name=model["name"],
+                    parameters=params,
+                    is_default_run=True,
+                    created_at = current_milli_time())
+
+    create_run(run)
+
+    # Store model ID to `tests` index with `status` set to `running`
+    body = {"status": "running", "created_at": run.created_at, "run_id": run.id}
+    es.index(index="tests", body=body, id=model_id)
+
+    return run_id
+
+
 def plugin_action(action_name, **kwargs):
     from src.settings import settings
 
diff --git a/api/validation/RunSchema.py b/api/validation/RunSchema.py
index 98a3330..18a81ea 100644
--- a/api/validation/RunSchema.py
+++ b/api/validation/RunSchema.py
@@ -95,4 +95,9 @@ class Config:
         ...,
         description="tasks",
         title="tasks",
-    )
\ No newline at end of file
+    )
+
+class RunStatusSchema(Enum):
+    success = "success"
+    running = "running"
+    failed = "failed"
\ No newline at end of file