# coding=utf-8 import json 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_save_connection_with_basic_auth(_fresh_store): connections.save_connection( name="secure", master="yarn", auth_type="basic", auth_user="hdfs", auth_password="secret", ) out = connections.get_connection("secure") assert out["auth_type"] == "basic" assert out["auth_user"] == "hdfs" assert out["auth_password"] == "secret" def test_save_connection_invalid_auth_type_raises(_fresh_store): with pytest.raises(ValueError, match="auth_type must be one of none/simple/basic/kerberos"): connections.save_connection( name="bad", master="yarn", auth_type="oauth", ) 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") def test_update_connection_changes_specified_field(_fresh_store): connections.save_connection(name="prod", master="yarn") out = connections.update_connection(name="prod", master="spark://new:7077") assert out["master"] == "spark://new:7077" assert _fresh_store.get("prod").master == "spark://new:7077" def test_update_connection_keeps_omitted_fields(_fresh_store): connections.save_connection( name="prod", master="yarn", deploy_mode="client", yarn_rm_url="http://rm:8088", ) out = connections.update_connection(name="prod", deploy_mode="cluster") assert out["deploy_mode"] == "cluster" assert out["master"] == "yarn" assert out["yarn_rm_url"] == "http://rm:8088" def test_update_connection_with_no_fields_is_noop(_fresh_store): connections.save_connection(name="prod", master="yarn", deploy_mode="client") out = connections.update_connection(name="prod") assert out["master"] == "yarn" assert out["deploy_mode"] == "client" def test_update_connection_allows_url_hosts_change(_fresh_store): connections.save_connection(name="prod", master="yarn") out = connections.update_connection(name="prod", allowed_url_hosts=["ccam*"]) assert out["allowed_url_hosts"] == ["ccam*"] def test_update_connection_replaces_not_merges_dict(_fresh_store): connections.save_connection(name="prod", master="yarn", spark_conf={"a": "1"}) out = connections.update_connection(name="prod", spark_conf={"b": "2"}) assert out["spark_conf"] == {"b": "2"} def test_update_connection_replaces_not_merges_list(_fresh_store): connections.save_connection( name="prod", master="yarn", allowed_url_hosts=["ccam*"], ) out = connections.update_connection(name="prod", allowed_url_hosts=["nm*"]) assert out["allowed_url_hosts"] == ["nm*"] def test_update_connection_raises_for_unknown_name(_fresh_store): with pytest.raises(KeyError, match="Connection not found"): connections.update_connection(name="missing", master="yarn") def test_update_connection_validates_patched_connection(_fresh_store): connections.save_connection(name="prod", master="yarn") with pytest.raises(ValueError, match="master must be"): connections.update_connection(name="prod", master="http://bad") def test_update_connection_persists_to_disk(_fresh_store, tmp_path): connections.save_connection(name="prod", master="yarn") connections.update_connection(name="prod", master="spark://new:7077") raw = json.loads((tmp_path / "connections.json").read_text()) assert raw["prod"]["master"] == "spark://new:7077"