diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d473b71 --- /dev/null +++ b/.gitignore @@ -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/ diff --git a/pyproject.toml b/pyproject.toml index 4b750d8..5a2110a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 diff --git a/spark_executor/core/__init__.py b/spark_executor/core/__init__.py new file mode 100644 index 0000000..36f7e1b --- /dev/null +++ b/spark_executor/core/__init__.py @@ -0,0 +1,5 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" diff --git a/spark_executor/models.py b/spark_executor/models.py new file mode 100644 index 0000000..17ba1cb --- /dev/null +++ b/spark_executor/models.py @@ -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 diff --git a/spark_executor/tools/__init__.py b/spark_executor/tools/__init__.py new file mode 100644 index 0000000..36f7e1b --- /dev/null +++ b/spark_executor/tools/__init__.py @@ -0,0 +1,5 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..36f7e1b --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,5 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..e0591c4 --- /dev/null +++ b/tests/conftest.py @@ -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)) diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..36f7e1b --- /dev/null +++ b/tests/unit/__init__.py @@ -0,0 +1,5 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py new file mode 100644 index 0000000..62e6e75 --- /dev/null +++ b/tests/unit/test_models.py @@ -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" diff --git a/uv.lock b/uv.lock index 615a86b..017e821 100644 --- a/uv.lock +++ b/uv.lock @@ -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"