chore: add pytest dev-deps and pydantic models (Job, Connection, PendingSubmission)

2026-06-24 14:34:09 +08:00
parent b56d7c80dd
commit 864093f1fb
10 changed files with 265 additions and 0 deletions
@@ -0,0 +1,22 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+.eggs/
+build/
+dist/
+
+# Virtual env
+.venv/
+
+# IDE
+.idea/
+.vscode/
+
+# Data directory (persisted Spark connections / pending jobs)
+data/
+
+# Test / coverage
+.pytest_cache/
+.coverage
+htmlcov/
@@ -12,6 +12,12 @@ dependencies = [
    "uvicorn>=0.49.0",
 ]

+[dependency-groups]
+dev = [
+    "pytest>=8.0",
+    "pytest-mock>=3.14",
+]
+
 [[tool.uv.index]]
 url = "https://pypi.tuna.tsinghua.edu.cn/simple/"
 default = true
@@ -0,0 +1,5 @@
+# coding=utf-8
+"""
+@Time   :2026/6/24
+@Author :tao.chen
+"""
@@ -0,0 +1,55 @@
+# coding=utf-8
+"""
+@Time   :2026/6/24
+@Author :tao.chen
+"""
+from datetime import datetime
+
+from pydantic import BaseModel, Field
+
+
+class Job(BaseModel):
+    job_id: str
+    application_id: str
+    script_path: str
+    queue: str
+    submit_time: datetime
+    connection: str
+
+
+class JobStatus(BaseModel):
+    application_id: str
+    state: str
+    raw: str = Field(default="")
+
+
+class SubmitResult(BaseModel):
+    job_id: str
+    application_id: str
+    tracking_url: str | None = None
+
+
+class Connection(BaseModel):
+    name: str
+    master: str
+    deploy_mode: str = "cluster"
+    yarn_rm_url: str | None = None
+    spark_conf: dict[str, str] = Field(default_factory=dict)
+
+
+class PendingSubmission(BaseModel):
+    pending_id: str
+    connection: str
+    master: str
+    deploy_mode: str
+    script_path: str
+    queue: str
+    executor_memory: str
+    executor_cores: int
+    num_executors: int
+    spark_conf: dict[str, str] = Field(default_factory=dict)
+    created_at: datetime
+    status: str = "PENDING"  # PENDING | SUBMITTED | CANCELLED | FAILED
+    error: str | None = None
+    job_id: str | None = None
+    application_id: str | None = None
@@ -0,0 +1,5 @@
+# coding=utf-8
+"""
+@Time   :2026/6/24
+@Author :tao.chen
+"""
@@ -0,0 +1,5 @@
+# coding=utf-8
+"""
+@Time   :2026/6/24
+@Author :tao.chen
+"""
@@ -0,0 +1,8 @@
+# coding=utf-8
+# Ensures the project root is on sys.path when running pytest from any cwd.
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
@@ -0,0 +1,5 @@
+# coding=utf-8
+"""
+@Time   :2026/6/24
+@Author :tao.chen
+"""
@@ -0,0 +1,87 @@
+# coding=utf-8
+from datetime import datetime
+from spark_executor.models import Connection, Job, JobStatus, PendingSubmission, SubmitResult
+
+
+def test_job_roundtrip():
+    job = Job(
+        job_id="abc123",
+        application_id="application_17400000001",
+        script_path="/tmp/jobs/job_001.py",
+        queue="default",
+        submit_time=datetime(2026, 6, 24, 10, 0, 0),
+        connection="prod-yarn",
+    )
+    dumped = job.model_dump()
+    assert dumped["job_id"] == "abc123"
+    assert dumped["application_id"] == "application_17400000001"
+    assert dumped["connection"] == "prod-yarn"
+
+
+def test_job_status_default_raw():
+    s = JobStatus(application_id="application_1", state="RUNNING")
+    assert s.raw == ""
+
+
+def test_submit_result_tracking_url_optional():
+    r = SubmitResult(job_id="j1", application_id="application_1", tracking_url=None)
+    assert r.tracking_url is None
+
+
+def test_connection_defaults():
+    c = Connection(name="prod", master="yarn")
+    assert c.deploy_mode == "cluster"
+    assert c.yarn_rm_url is None
+    assert c.spark_conf == {}
+
+
+def test_connection_with_spark_conf():
+    c = Connection(
+        name="dev",
+        master="spark://master:7077",
+        deploy_mode="client",
+        spark_conf={"spark.sql.shuffle.partitions": "200"},
+    )
+    assert c.spark_conf["spark.sql.shuffle.partitions"] == "200"
+
+
+def test_pending_submission_defaults_to_pending_status():
+    p = PendingSubmission(
+        pending_id="p_1",
+        connection="prod",
+        master="yarn",
+        deploy_mode="cluster",
+        script_path="/tmp/j.py",
+        queue="default",
+        executor_memory="4G",
+        executor_cores=2,
+        num_executors=2,
+        spark_conf={},
+        created_at=datetime(2026, 6, 24),
+    )
+    assert p.status == "PENDING"
+    assert p.error is None
+    assert p.job_id is None
+    assert p.application_id is None
+
+
+def test_pending_submission_can_record_outcome():
+    p = PendingSubmission(
+        pending_id="p_2",
+        connection="prod",
+        master="yarn",
+        deploy_mode="cluster",
+        script_path="/tmp/j.py",
+        queue="default",
+        executor_memory="4G",
+        executor_cores=2,
+        num_executors=2,
+        spark_conf={},
+        created_at=datetime(2026, 6, 24),
+        status="SUBMITTED",
+        job_id="j_abc",
+        application_id="application_17400000001",
+    )
+    assert p.status == "SUBMITTED"
+    assert p.job_id == "j_abc"
+    assert p.application_id == "application_17400000001"
@@ -349,6 +349,15 @@ wheels = [
    { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455 },
 ]

+[[package]]
+name = "iniconfig"
+version = "2.3.0"
+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503 }
+wheels = [
+    { url = "https://pypi.tuna.tsinghua.edu.cn/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484 },
+]
+
 [[package]]
 name = "jsonschema"
 version = "4.26.0"
@@ -435,6 +444,24 @@ wheels = [
    { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979 },
 ]

+[[package]]
+name = "packaging"
+version = "26.2"
+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134 }
+wheels = [
+    { url = "https://pypi.tuna.tsinghua.edu.cn/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195 },
+]
+
+[[package]]
+name = "pluggy"
+version = "1.6.0"
+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412 }
+wheels = [
+    { url = "https://pypi.tuna.tsinghua.edu.cn/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538 },
+]
+
 [[package]]
 name = "pycparser"
 version = "3.0"
@@ -567,6 +594,34 @@ crypto = [
    { name = "cryptography" },
 ]

+[[package]]
+name = "pytest"
+version = "9.1.1"
+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "iniconfig" },
+    { name = "packaging" },
+    { name = "pluggy" },
+    { name = "pygments" },
+]
+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369 }
+wheels = [
+    { url = "https://pypi.tuna.tsinghua.edu.cn/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536 },
+]
+
+[[package]]
+name = "pytest-mock"
+version = "3.15.1"
+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
+dependencies = [
+    { name = "pytest" },
+]
+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/68/14/eb014d26be205d38ad5ad20d9a80f7d201472e08167f0bb4361e251084a9/pytest_mock-3.15.1.tar.gz", hash = "sha256:1849a238f6f396da19762269de72cb1814ab44416fa73a8686deac10b0d87a0f", size = 34036 }
+wheels = [
+    { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl", hash = "sha256:0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d", size = 10095 },
+]
+
 [[package]]
 name = "python-dotenv"
 version = "1.2.2"
@@ -779,6 +834,12 @@ dependencies = [
    { name = "uvicorn" },
 ]

+[package.dev-dependencies]
+dev = [
+    { name = "pytest" },
+    { name = "pytest-mock" },
+]
+
 [package.metadata]
 requires-dist = [
    { name = "fastapi", specifier = ">=0.138.0" },
@@ -790,6 +851,12 @@ requires-dist = [
    { name = "uvicorn", specifier = ">=0.49.0" },
 ]

+[package.metadata.requires-dev]
+dev = [
+    { name = "pytest", specifier = ">=8.0" },
+    { name = "pytest-mock", specifier = ">=3.14" },
+]
+
 [[package]]
 name = "sse-starlette"
 version = "3.4.5"