Modify upon review

Signed-off-by: Alina Buzachis <[email protected]>

Author: alinabuzachis

core/task_runner.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+from contextlib import closing
 
 from core.utils.controller import assign_execute_roles
 from core.utils.controller import build_collection_uri
@@ -9,6 +10,7 @@
 from core.utils.controller import create_labels
 from core.utils.controller import create_project
 from core.utils.controller import download_collection
+from core.utils.controller import get_http_session
 from core.utils.controller import save_instance_state
 
 from .models import Pattern
@@ -74,19 +76,23 @@ def run_pattern_instance_task(instance_id: int, task_id: int) -> None:
         if not pattern_def:
             raise ValueError("Pattern definition is missing.")
 
-        task.mark_running({"info": "Creating controller project"})
-        project_id = create_project(instance, pattern, pattern_def)
-        task.mark_running({"info": "Creating execution environment"})
-        ee_id = create_execution_environment(instance, pattern_def)
-        task.mark_running({"info": "Creating labels"})
-        labels = create_labels(instance, pattern_def)
-        task.mark_running({"info": "Creating job templates"})
-        automations = create_job_templates(instance, pattern_def, project_id, ee_id)
-        task.mark_running({"info": "Saving instance"})
-        save_instance_state(instance, project_id, ee_id, labels, automations)
-        task.mark_running({"info": "Assigning roles"})
-        assign_execute_roles(instance.executors, automations)
-        task.mark_completed({"info": "PatternInstance processed"})
+        # Create a single session for all AAP calls
+        with closing(get_http_session()) as session:
+            task.mark_running({"info": "Creating controller project"})
+            project_id = create_project(session, instance, pattern)
+            task.mark_running({"info": "Creating execution environment"})
+            ee_id = create_execution_environment(session, instance, pattern_def)
+            task.mark_running({"info": "Creating labels"})
+            labels = create_labels(session, instance, pattern_def)
+            task.mark_running({"info": "Creating job templates"})
+            automations = create_job_templates(
+                session, instance, pattern_def, project_id, ee_id
+            )
+            task.mark_running({"info": "Saving instance"})
+            save_instance_state(instance, project_id, ee_id, labels, automations)
+            task.mark_running({"info": "Assigning roles"})
+            assign_execute_roles(session, instance.executors, automations)
+            task.mark_completed({"info": "PatternInstance processed"})
     except Exception as e:
         logger.exception("Failed to process PatternInstance.")
         task.mark_failed({"error": str(e)})

core/tests/test_controller_client.py

@@ -34,7 +34,7 @@ def test_post_success(mock_get_http_session):
     session.post.return_value = _fake_response(201, {"id": 99})
     mock_get_http_session.return_value = session
 
-    assert cc.post("/labels/", {"name": "foo"}) == {"id": 99}
+    assert cc.post(session, "/labels/", {"name": "foo"}) == {"id": 99}
     session.post.assert_called_once()
 
 
@@ -47,7 +47,7 @@ def test_post_duplicate_name(mock_get_http_session):
     )
     mock_get_http_session.return_value = session
 
-    out = cc.post("/labels/", {"name": "foo"})
+    out = cc.post(session, "/labels/", {"name": "foo"})
     assert out == {"id": 42, "name": "foo"}
 
     session.get.assert_called_once()
@@ -61,7 +61,7 @@ def test_post_non_400_error_is_propagated(mock_get_http_session):
     mock_get_http_session.return_value = session
 
     with pytest.raises(requests.HTTPError):
-        cc.post("/labels/", {"name": "foo"})
+        cc.post(session, "/labels/", {"name": "foo"})
 
     # Ensure we don't try to deduplicate
     session.get.assert_not_called()
@@ -76,7 +76,7 @@ def test_post_400_dedupe_no_results_raises(mock_get_http_session):
     mock_get_http_session.return_value = session
 
     with pytest.raises(requests.HTTPError) as exc:
-        cc.post("/labels/", {"name": "foo", "organization": 1})
+        cc.post(session, "/labels/", {"name": "foo", "organization": 1})
 
     assert isinstance(exc.value, requests.HTTPError)
     assert exc.value.response.status_code == 400
@@ -95,7 +95,7 @@ def test_post_400_dedupe_lookup_http_error_still_raises(mock_get_http_session):
     mock_get_http_session.return_value = session
 
     with pytest.raises(requests.HTTPError):
-        cc.post("/labels/", {"name": "foo"})
+        cc.post(session, "/labels/", {"name": "foo"})
 
     session.get.assert_called_once()
 
@@ -108,7 +108,7 @@ def test_post_400_dedupe_params_subset_when_missing_org(mock_get_http_session):
     mock_get_http_session.return_value = session
 
     with pytest.raises(requests.HTTPError):
-        cc.post("/labels/", {"name": "foo"})
+        cc.post(session, "/labels/", {"name": "foo"})
 
     # Only 'name' present in params, no 'organization' since it wasn’t provided
     _, kwargs = session.get.call_args
@@ -124,6 +124,7 @@ def test_post_400_with_custom_dedupe_keys(mock_get_http_session):
 
     with pytest.raises(requests.HTTPError):
         cc.post(
+            session,
             "/labels/",
             {"name": "foo", "organization": 1, "extra": "x"},
             dedupe_keys=("name", "extra"),
@@ -148,7 +149,7 @@ def test_post_400_error_json_fallback_to_text(mock_get_http_session):
     mock_get_http_session.return_value = session
 
     with pytest.raises(requests.HTTPError) as exc:
-        cc.post("/labels/", {"name": "foo"})
+        cc.post(session, "/labels/", {"name": "foo"})
 
     # Error message should include the fallback text
     assert "not json" in str(exc.value)