From a0be90ae53716e7b17beb3ace241bd2168f29f4a Mon Sep 17 00:00:00 2001
From: karanjot786 <karanjots801@gmail.com>
Date: Sat, 12 Oct 2024 21:22:37 +0530
Subject: [PATCH] feat: add TES models with unit tests

---
 crategen/models/__init__.py   | 32 ++++++++++++++++++++++++++++++++
 crategen/models/tes_models.py | 26 +++++++++++++++++++++++++-
 tests/unit/test_tes_models.py | 22 +++++++++++++---------
 3 files changed, 70 insertions(+), 10 deletions(-)

diff --git a/crategen/models/__init__.py b/crategen/models/__init__.py
index e69de29..2da85d7 100644
--- a/crategen/models/__init__.py
+++ b/crategen/models/__init__.py
@@ -0,0 +1,32 @@
+"""Models package for TES, WES, and WRROC.
+
+This package contains Pydantic models that conform to the GA4GH schemas for Task Execution Services (TES),
+Workflow Execution Services (WES), and WRROC. These models are used for data validation and type safety
+throughout the CrateGen project.
+"""
+
+from .tes_models import (
+    TESData,
+    TESExecutor,
+    TESExecutorLog,
+    TESFileType,
+    TESInput,
+    TESOutput,
+    TESOutputFileLog,
+    TESResources,
+    TESState,
+    TESTaskLog,
+)
+
+__all__ = [
+    "TESData",
+    "TESInput",
+    "TESOutput",
+    "TESExecutor",
+    "TESTaskLog",
+    "TESResources",
+    "TESExecutorLog",
+    "TESOutputFileLog",
+    "TESFileType",
+    "TESState",
+]
diff --git a/crategen/models/tes_models.py b/crategen/models/tes_models.py
index 1f64011..ce4efda 100644
--- a/crategen/models/tes_models.py
+++ b/crategen/models/tes_models.py
@@ -11,11 +11,31 @@
 
 
 class TESFileType(str, Enum):
+    """Enumeration of TES file types.
+
+    **Attributes:**
+
+    - **FILE** (`str`): Represents a file.
+    - **DIRECTORY** (`str`): Represents a directory.
+    """
     FILE = "FILE"
     DIRECTORY = "DIRECTORY"
 
 
 class TESState(str, Enum):
+    """Enumeration of TES task states.
+
+    **Attributes:**
+    - **UNKNOWN** (`str`): The task state is unknown.
+    - **QUEUED** (`str`): The task is queued.
+    - **INITIALIZING** (`str`): The task is initializing.
+    - **RUNNING** (`str`): The task is running.
+    - **PAUSED** (`str`): The task is paused.
+    - **COMPLETE** (`str`): The task is complete.
+    - **EXECUTOR_ERROR** (`str`): The task encountered an executor error.
+    - **SYSTEM_ERROR** (`str`): The task encountered a system error.
+    - **CANCELLED** (`str`): The task was cancelled.
+    """
     UNKNOWN = "UNKNOWN"
     QUEUED = "QUEUED"
     INITIALIZING = "INITIALIZING"
@@ -66,11 +86,12 @@ class TESExecutorLog(BaseModel):
 
     @validator("start_time", "end_time", pre=True, always=True)
     def validate_datetime(cls, value):
+        """Convert start and end times to RFC 3339 format."""
         return convert_to_iso8601(value)
 
 
 class TESExecutor(BaseModel):
-    """An array of executors to be run
+    """An array of executors to be run.
 
     **Attributes:**
     - **image** (`str`): Name of the container image.
@@ -164,6 +185,7 @@ def validate_content_and_url(cls, values):
 
     @validator("path")
     def validate_path(cls, value):
+        """Validate that the path is an absolute path."""
         if not os.path.isabs(value):
             raise ValueError("The 'path' attribute must contain an absolute path.")
         return value
@@ -191,6 +213,7 @@ class TESOutput(BaseModel):
 
     @validator("path")
     def validate_path(cls, value):
+        """Validate that the path is an absolute path."""
         if not os.path.isabs(value):
             raise ValueError("The 'path' attribute must contain an absolute path.")
         return value
@@ -221,6 +244,7 @@ class TESTaskLog(BaseModel):
 
     @validator("start_time", "end_time", pre=True, always=True)
     def validate_datetime(cls, value):
+        """Convert start and end times to RFC 3339 format."""
         return convert_to_iso8601(value)
 
 
diff --git a/tests/unit/test_tes_models.py b/tests/unit/test_tes_models.py
index ef5f1c9..eb06ad4 100644
--- a/tests/unit/test_tes_models.py
+++ b/tests/unit/test_tes_models.py
@@ -6,15 +6,19 @@
 from pydantic import ValidationError
 
 from crategen.models.tes_models import (
-    TESInput,
-    TESOutput,
-    TESExecutor,
     TESData,
+    TESExecutor,
     TESFileType,
+    TESInput,
+    TESOutput,
+    TESResources,
     TESState,
-    TESResources
 )
 
+EXPECTED_CPU_CORES = 2
+EXPECTED_RAM_GB = 4.0
+EXPECTED_DISK_GB = 10.0
+EXPECTED_PREEMPTIBLE = False
 
 def test_tes_input_with_url():
     """Test TESInput model with a valid URL and absolute path."""
@@ -174,10 +178,10 @@ def test_tes_data_valid():
             )
         ],
         resources=TESResources(
-            cpu_cores=2,
-            ram_gb=4.0,
-            disk_gb=10.0,
-            preemptible=False,
+            cpu_cores=EXPECTED_CPU_CORES,
+            ram_gb=EXPECTED_RAM_GB,
+            disk_gb=EXPECTED_DISK_GB,
+            preemptible=EXPECTED_PREEMPTIBLE,
         ),
         volumes=["/data"],
         tags={"project": "CrateGen"},
@@ -186,7 +190,7 @@ def test_tes_data_valid():
     assert tes_data.inputs[0].url == "https://raw.githubusercontent.com/elixir-cloud-aai/CrateGen/refs/heads/main/README.md"
     assert tes_data.outputs[0].path == "/data/output/LICENSE"
     assert tes_data.executors[0].image == "python:3.8-slim"
-    assert tes_data.resources.cpu_cores == 2
+    assert tes_data.resources.cpu_cores == EXPECTED_CPU_CORES
     assert tes_data.volumes == ["/data"]
     assert tes_data.tags == {"project": "CrateGen"}