agent-app / tests /test_aml_submit.py
ganeshkota's picture
Add AML pipeline mode, sanitize configs, and production docs
9a45d9a
Raw
History Blame Contribute Delete
2.63 kB
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
sys.path.insert(0, str(SRC))
from agent_app.revenue_assurance.aml_submit import AmlSubmitConfig, build_job_spec, load_submit_config
class AmlSubmitTests(unittest.TestCase):
def test_build_job_spec_uses_config_path_and_defaults(self) -> None:
with tempfile.TemporaryDirectory() as tmpdir, patch.dict(
os.environ,
{
"AML_SUBSCRIPTION_ID": "",
"AML_RESOURCE_GROUP": "",
"AML_WORKSPACE_NAME": "",
"AML_COMPUTE_NAME": "",
},
clear=False,
):
model_config_path = Path(tmpdir) / "model_config.json"
model_config_path.write_text(
json.dumps(
{
"model_endpoint": "https://example.services.ai.azure.com/openai/v1/responses",
"model_deployment": "gpt-5.6-sol",
"api_version": "2026-07-09",
"api_key": "secret",
"use_agentic": True,
}
)
)
submit_config_path = Path(tmpdir) / "submit_config.json"
submit_config_path.write_text(
json.dumps(
{
"subscription_id": "sub-id",
"resource_group": "rg-name",
"workspace_name": "ws-name",
"compute_name": "cpu-cluster",
"code_path": ".",
"environment_image": "python:3.12-slim",
"experiment_name": "revenue-assurance",
"display_name": "revenue-assurance-agent",
"revenue_assurance_config_path": str(model_config_path),
}
)
)
submit_config = load_submit_config(str(submit_config_path))
self.assertIsInstance(submit_config, AmlSubmitConfig)
spec = build_job_spec(submit_config=submit_config)
self.assertEqual(spec["compute"], "cpu-cluster")
self.assertTrue(spec["environment_variables"]["REVENUE_ASSURANCE_USE_AGENTIC"])
self.assertIn("python -m agent_app.revenue_assurance.aml_runner", spec["command"])
self.assertEqual(spec["model_deployment"], "gpt-5.6-sol")
if __name__ == "__main__":
unittest.main()