Lets try this...

Author: jmchilton

pulsar/client/client.py

@@ -3,6 +3,7 @@
 from enum import Enum
 from typing import (
     Any,
+    cast,
     Callable,
     Dict,
     Optional,
@@ -1110,7 +1111,7 @@ def raw_check_complete(self) -> Dict[str, Any]:
 
 
 def gcp_state_to_pulsar_status(state: Optional["batch_v1.JobStatus.State"]) -> str:
-    state = state or batch_v1.JobStatus.State.STATE_UNSPECIFIED
+    state = state or cast(batch_v1.JobStatus.State, batch_v1.JobStatus.State.STATE_UNSPECIFIED)
     # STATE_UNSPECIFIED	Job state unspecified.
     # QUEUED	Job is admitted (validated and persisted) and waiting for resources.
     # SCHEDULED	Job is scheduled to run as soon as resource allocation is ready. The resource
@@ -1141,7 +1142,7 @@ def gcp_state_to_pulsar_status(state: Optional["batch_v1.JobStatus.State"]) -> s
 
 
 def gcp_state_is_complete(state: Optional["batch_v1.JobStatus.State"]) -> bool:
-    state = state or batch_v1.JobStatus.State.STATE_UNSPECIFIED
+    state = state or cast(batch_v1.JobStatus.State, batch_v1.JobStatus.State.STATE_UNSPECIFIED)
     state_map = {
         batch_v1.JobStatus.State.STATE_UNSPECIFIED: True,
         batch_v1.JobStatus.State.QUEUED: False,

pulsar/client/container_job_config.py

@@ -12,7 +12,6 @@
 from typing import (
     Dict,
     List,
-    Literal,
     NamedTuple,
     Optional,
 )
@@ -22,6 +21,8 @@
     Field,
 )
 from pydantictes.models import TesResources
+from typing_extensions import Literal
+
 from pulsar.managers.util.gcp_util import (
     batch_v1,
     ensure_client as ensure_gcp_client,
@@ -197,7 +198,7 @@ def gcp_job_template(params: GcpJobParams) -> "batch_v1.Job":
     job.labels = params.labels or {}
     # We use Cloud Logging as it's an out of the box available option
     job.logs_policy = batch_v1.LogsPolicy()
-    job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING
+    job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING  # type: ignore[assignment]
 
     return job
 
@@ -230,19 +231,12 @@ class BasicAuth(BaseModel):
     password: str = Field(..., description="Password for basic authentication.")
 
 
-class TesJobParams(BaseModel):
+class TesJobParams(TesResources):
     tes_url: str = Field(..., description="URL of the TES service.")
     authorization: Literal["none", "basic"] = Field(
         "none", description="Authorization type for TES service."
     )
     basic_auth: Optional[BasicAuth] = Field(None, description="Authorization for TES service.")
-    cpu_cores: Optional[int] = TesResources.__pydantic_fields__["cpu_cores"]
-    preemptible: Optional[bool] = TesResources.__pydantic_fields__["preemptible"]
-    ram_gb: Optional[float] = TesResources.__pydantic_fields__["ram_gb"]
-    disk_gb: Optional[float] = TesResources.__pydantic_fields__["disk_gb"]
-    zones: Optional[List[str]] = TesResources.__pydantic_fields__["zones"]
-    backend_parameters: Optional[Dict[str, str]] = TesResources.__pydantic_fields__["backend_parameters"]
-    backend_parameters_strict: Optional[bool] = TesResources.__pydantic_fields__["backend_parameters_strict"]
 
 
 def parse_tes_job_params(params: dict) -> TesJobParams:
@@ -290,19 +284,5 @@ def tes_client_from_params(tes_params: TesJobParams) -> TesClient:
 
 
 def tes_resources(tes_params: TesJobParams) -> TesResources:
-    cpu_cores: Optional[int] = tes_params.cpu_cores
-    preemptible: Optional[bool] = tes_params.preemptible
-    ram_gb: Optional[float] = tes_params.ram_gb
-    disk_gb: Optional[float] = tes_params.disk_gb
-    zones: Optional[List[str]] = tes_params.zones
-    backend_parameters: Optional[Dict[str, str]] = tes_params.backend_parameters
-    backend_parameters_strict: Optional[bool] = tes_params.backend_parameters_strict
-    return TesResources(
-        cpu_cores=cpu_cores,
-        preemptible=preemptible,
-        ram_gb=ram_gb,
-        disk_gb=disk_gb,
-        zones=zones,
-        backend_parameters=backend_parameters,
-        backend_parameters_strict=backend_parameters_strict,
-    )
+    # TesJobParams subclasses it so just pass through as is.
+    return tes_params

pulsar/managers/util/gcp_util.py

@@ -8,8 +8,8 @@
     from google.cloud import batch_v1  # type: ignore
     from google.oauth2 import service_account  # type: ignore
 except ImportError as exc:
-    service_account = None
-    batch_v1 = None
+    service_account = None  # type: ignore[assignment]
+    batch_v1 = None  # type: ignore[assignment]
     GCP_IMPORT_MESSAGE = (
         "The Python google-cloud-batch package is required to use "
         "this feature, please install it or correct the "

test/tes_test.py

@@ -1,14 +1,18 @@
-from pulsar.managers.util.tes import tes_resources
+from pulsar.client.container_job_config import (
+    parse_tes_job_params,
+    tes_resources,
+)
 
 
 def test_tes_resources_from_xml():
-    resources = tes_resources({
+    resources = tes_resources(parse_tes_job_params({
+        "tes_url": "http://moo",
         "tes_cpu_cores": "2",
         "tes_preemptible": "true",
         "tes_ram_gb": "128.0",
         "tes_disk_gb": "512.0",
         "tes_zones": "us-west-1,us-east-1",
-    })
+    }))
     assert resources.cpu_cores == 2
     assert resources.preemptible is True
     assert resources.ram_gb == 128.0
@@ -19,13 +23,14 @@ def test_tes_resources_from_xml():
 
 
 def test_tes_resources_from_yaml():
-    resources = tes_resources({
+    resources = tes_resources(parse_tes_job_params({
+        "tes_url": "http://moo",
         "tes_cpu_cores": 4,
         "tes_ram_gb": 127.0,
         "tes_disk_gb": 513.0,
         "tes_zones": ["us-west-1", "us-east-1"],
         "tes_backend_parameters": {"moo": "cow"},
-    })
+    }))
     assert resources.cpu_cores == 4
     assert resources.preemptible is None
     assert resources.ram_gb == 127.0