# coding=utf-8 from pathlib import Path import pytest from spark_executor.core import connection_store from spark_executor.tools import connections @pytest.fixture(autouse=True) def _fresh_store(tmp_path: Path, monkeypatch): monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) monkeypatch.setattr(connections, "store", connection_store.store) return connection_store.store # --- save_connection --- def test_save_connection(_fresh_store): out = connections.save_connection(name="prod", master="yarn") assert out == {"name": "prod", "status": "SAVED"} assert _fresh_store.get("prod") is not None def test_save_connection_upserts(_fresh_store): connections.save_connection(name="prod", master="yarn") connections.save_connection(name="prod", master="spark://new:7077", deploy_mode="client") assert _fresh_store.get("prod").master == "spark://new:7077" assert _fresh_store.get("prod").deploy_mode == "client" def test_save_connection_with_spark_conf(_fresh_store): connections.save_connection( name="dev", master="yarn", spark_conf={"spark.sql.shuffle.partitions": "200"}, ) assert _fresh_store.get("dev").spark_conf["spark.sql.shuffle.partitions"] == "200" # --- list_connections --- def test_list_connections_empty(_fresh_store): assert connections.list_connections() == [] def test_list_connections_returns_all_saved(_fresh_store): connections.save_connection(name="a", master="yarn") connections.save_connection(name="b", master="spark://m:7077", deploy_mode="client") out = connections.list_connections() names = {c["name"] for c in out} assert names == {"a", "b"} masters = {c["name"]: c["master"] for c in out} assert masters == {"a": "yarn", "b": "spark://m:7077"} # --- get_connection --- def test_get_connection_returns_dict(_fresh_store): connections.save_connection(name="prod", master="yarn", deploy_mode="cluster") out = connections.get_connection("prod") assert out["name"] == "prod" assert out["master"] == "yarn" assert out["deploy_mode"] == "cluster" assert out["yarn_rm_url"] is None assert out["spark_conf"] == {} assert out["ssl_verify"] is None assert out["ssl_ca_bundle"] is None def test_save_connection_preserves_ssl_fields(_fresh_store): connections.save_connection( name="secure", master="yarn", ssl_verify=False, ssl_ca_bundle="/tmp/ca.crt", ) out = connections.get_connection("secure") assert out["ssl_verify"] is False assert out["ssl_ca_bundle"] == "/tmp/ca.crt" def test_get_connection_unknown_raises(_fresh_store): with pytest.raises(KeyError): connections.get_connection("missing") # --- delete_connection --- def test_delete_connection_returns_status(_fresh_store): connections.save_connection(name="prod", master="yarn") assert connections.delete_connection("prod") == {"name": "prod", "status": "DELETED"} assert _fresh_store.get("prod") is None def test_delete_connection_unknown_raises(_fresh_store): with pytest.raises(KeyError): connections.delete_connection("missing")