From 9cbc177c55e6e396d309bba323b0c00ca2e2f0ff Mon Sep 17 00:00:00 2001 From: Michael Abashian Date: Tue, 21 Jan 2025 16:59:03 -0500 Subject: [PATCH 1/4] Adds tests for role generation anonymization --- .../ai/api/tests/test_role_generation_view.py | 140 ++++++++++++++++++ ansible_ai_connect/ai/api/tests/test_views.py | 33 ----- ansible_ai_connect/ai/api/views.py | 5 + 3 files changed, 145 insertions(+), 33 deletions(-) create mode 100644 ansible_ai_connect/ai/api/tests/test_role_generation_view.py diff --git a/ansible_ai_connect/ai/api/tests/test_role_generation_view.py b/ansible_ai_connect/ai/api/tests/test_role_generation_view.py new file mode 100644 index 000000000..fe522bac1 --- /dev/null +++ b/ansible_ai_connect/ai/api/tests/test_role_generation_view.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 + +# Copyright Red Hat +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import logging +import uuid +from http import HTTPStatus +from typing import Optional +from unittest.mock import Mock, patch + +from django.apps import apps +from django.test import override_settings +from django.urls import reverse + +from ansible_ai_connect.ai.api.model_pipelines.config_pipelines import BaseConfig +from ansible_ai_connect.ai.api.model_pipelines.dummy.pipelines import ROLE_FILES +from ansible_ai_connect.ai.api.model_pipelines.pipelines import ( + ModelPipelineRoleGeneration, + RoleGenerationParameters, + RoleGenerationResponse, +) +from ansible_ai_connect.ai.api.model_pipelines.tests import mock_config +from ansible_ai_connect.healthcheck.backends import HealthCheckSummary +from ansible_ai_connect.test_utils import ( + WisdomAppsBackendMocking, + WisdomServiceAPITestCaseBase, +) + +logger = logging.getLogger(__name__) + + +class MockedConfig(BaseConfig): + def __init__(self): + super().__init__(inference_url="mock-url", model_id="mock-model", timeout=None) + + +class MockedPipelineRoleGeneration(ModelPipelineRoleGeneration[MockedConfig]): + + def __init__(self, response_roles: str, response_files: list, response_outline: str): + super().__init__(MockedConfig()) + self.response_roles = response_roles + self.response_files = response_files + self.response_outline = response_outline + + def invoke(self, params: RoleGenerationParameters) -> RoleGenerationResponse: + return self.response_roles, self.response_files, self.response_outline + + def self_test(self) -> Optional[HealthCheckSummary]: + raise NotImplementedError + + +@override_settings(ANSIBLE_AI_MODEL_MESH_CONFIG=mock_config("dummy")) +class TestRoleGenerationView(WisdomAppsBackendMocking, WisdomServiceAPITestCaseBase): + + @override_settings(SEGMENT_WRITE_KEY="DUMMY_KEY_VALUE") + def test_ok(self): + generation_id = uuid.uuid4() + payload = { + "text": "Set up MySQL and email the admin", + "generationId": generation_id, + "ansibleExtensionVersion": "24.4.0", + } + self.client.force_authenticate(user=self.user) + with self.assertLogs(logger="root", level="DEBUG") as log: + r = self.client.post(reverse("generations/role"), payload, format="json") + segment_events = self.extractSegmentEventsFromLog(log) + roleGenEvent = segment_events[0] + self.assertEqual(r.status_code, HTTPStatus.OK) + self.assertIsNotNone(r.data) + self.assertEqual(r.data["files"], ROLE_FILES) + self.assertEqual(r.data["format"], "plaintext") + self.assertEqual(r.data["generationId"], generation_id) + self.assertEqual(r.data["outline"], "") + self.assertEqual(r.data["role"], "install_nginx") + self.assertEqual(roleGenEvent["event"], "codegenRole") + self.assertEqual(roleGenEvent["properties"]["generationId"], str(generation_id)) + + def test_unauthorized(self): + payload = {} + # Hit the API without authentication + r = self.client.post(reverse("generations/role"), payload, format="json") + self.assertEqual(r.status_code, HTTPStatus.UNAUTHORIZED) + + @override_settings(ANSIBLE_AI_ENABLE_TECH_PREVIEW=True) + def test_with_anonymized_response(self): + generation_id = str(uuid.uuid4()) + payload = { + "text": "Install mysql and email admin@redhat.com", + "generationId": generation_id, + "ansibleExtensionVersion": "24.4.0", + } + + with patch.object( + apps.get_app_config("ai"), + "get_model_pipeline", + Mock( + return_value=MockedPipelineRoleGeneration( + "Install mysql and email admin@redhat.com", + [ + { + "name": "main.yml", + "content": """ + --- + - name: Install MySQL + package: + name: mysql + state: present + - name: Email admin + mail: + to: admin@redhat.com + """, + } + ], + "Install mysql and email admin@redhat.com", + ) + ), + ): + self.client.force_authenticate(user=self.user) + r = self.client.post(reverse("generations/role"), payload, format="json") + self.assertEqual(r.status_code, HTTPStatus.OK) + self.assertIsNotNone(r.data["role"]) + self.assertEqual(len(r.data["files"]), 1) + self.assertIsNotNone(r.data["outline"]) + self.assertTrue("mysql" in r.data["role"]) + self.assertTrue("mysql" in r.data["outline"]) + self.assertFalse("admin@redhat.com" in r.data["role"]) + self.assertFalse("admin@redhat.com" in r.data["outline"]) + for file in r.data["files"]: + self.assertFalse("admin@redhat.com" in file["content"]) diff --git a/ansible_ai_connect/ai/api/tests/test_views.py b/ansible_ai_connect/ai/api/tests/test_views.py index 1df09bd64..7ef29a81b 100644 --- a/ansible_ai_connect/ai/api/tests/test_views.py +++ b/ansible_ai_connect/ai/api/tests/test_views.py @@ -63,7 +63,6 @@ WcaValidationFailureException, ) from ansible_ai_connect.ai.api.model_pipelines.config_pipelines import BaseConfig -from ansible_ai_connect.ai.api.model_pipelines.dummy.pipelines import ROLE_FILES from ansible_ai_connect.ai.api.model_pipelines.exceptions import ( ModelTimeoutError, WcaBadRequest, @@ -3507,38 +3506,6 @@ def test_with_custom_prompt_missing_outline_when_not_needed(self): self.assertEqual(args.text, "Install nginx on RHEL9 isabella13@example.com") -@override_settings(ANSIBLE_AI_MODEL_MESH_CONFIG=mock_config("dummy")) -class TestRoleGenerationView(WisdomAppsBackendMocking, WisdomServiceAPITestCaseBase): - @override_settings(SEGMENT_WRITE_KEY="DUMMY_KEY_VALUE") - def test_ok(self): - generation_id = uuid.uuid4() - payload = { - "text": "Install nginx and enable the service", - "generationId": generation_id, - "ansibleExtensionVersion": "24.4.0", - } - self.client.force_authenticate(user=self.user) - with self.assertLogs(logger="root", level="DEBUG") as log: - r = self.client.post(reverse("generations/role"), payload, format="json") - segment_events = self.extractSegmentEventsFromLog(log) - roleGenEvent = segment_events[0] - self.assertEqual(r.status_code, HTTPStatus.OK) - self.assertIsNotNone(r.data) - self.assertEqual(r.data["files"], ROLE_FILES) - self.assertEqual(r.data["format"], "plaintext") - self.assertEqual(r.data["generationId"], generation_id) - self.assertEqual(r.data["outline"], "") - self.assertEqual(r.data["role"], "install_nginx") - self.assertEqual(roleGenEvent["event"], "codegenRole") - self.assertEqual(roleGenEvent["properties"]["generationId"], str(generation_id)) - - def test_unauthorized(self): - payload = {} - # Hit the API without authentication - r = self.client.post(reverse("generations/role"), payload, format="json") - self.assertEqual(r.status_code, HTTPStatus.UNAUTHORIZED) - - @override_settings(ANSIBLE_AI_MODEL_MESH_CONFIG=mock_config("wca")) @override_settings(ANSIBLE_AI_ENABLE_PLAYBOOK_ENDPOINT=True) class TestGenerationViewWithWCA(WisdomAppsBackendMocking, WisdomServiceAPITestCaseBase): diff --git a/ansible_ai_connect/ai/api/views.py b/ansible_ai_connect/ai/api/views.py index f04c7db00..e8920522f 100644 --- a/ansible_ai_connect/ai/api/views.py +++ b/ansible_ai_connect/ai/api/views.py @@ -925,6 +925,11 @@ def post(self, request) -> Response: outline, value_template=Template("{{ _${variable_name}_ }}") ) + for file in files: + file["content"] = anonymizer.anonymize_struct( + file["content"], value_template=Template("{{ _${variable_name}_ }}") + ) + answer = { "role": anonymized_role, "outline": anonymized_outline, From e86f8e6e701f5dad881789b947db4a9b29ea929d Mon Sep 17 00:00:00 2001 From: Michael Abashian Date: Wed, 22 Jan 2025 08:23:22 -0500 Subject: [PATCH 2/4] Remove unnecessary logic to anonymize role file content. This is already handled by anonymize_struct in the call above --- ansible_ai_connect/ai/api/views.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/ansible_ai_connect/ai/api/views.py b/ansible_ai_connect/ai/api/views.py index e8920522f..f04c7db00 100644 --- a/ansible_ai_connect/ai/api/views.py +++ b/ansible_ai_connect/ai/api/views.py @@ -925,11 +925,6 @@ def post(self, request) -> Response: outline, value_template=Template("{{ _${variable_name}_ }}") ) - for file in files: - file["content"] = anonymizer.anonymize_struct( - file["content"], value_template=Template("{{ _${variable_name}_ }}") - ) - answer = { "role": anonymized_role, "outline": anonymized_outline, From fa115b63df79e3cdf7c8fed31ef220b3e9b20f69 Mon Sep 17 00:00:00 2001 From: Michael Abashian Date: Wed, 22 Jan 2025 11:29:08 -0500 Subject: [PATCH 3/4] Re-add logic to anonymize files but pass the root object rather than iterating over each file --- ansible_ai_connect/ai/api/views.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ansible_ai_connect/ai/api/views.py b/ansible_ai_connect/ai/api/views.py index f04c7db00..d6dc976ec 100644 --- a/ansible_ai_connect/ai/api/views.py +++ b/ansible_ai_connect/ai/api/views.py @@ -924,11 +924,14 @@ def post(self, request) -> Response: anonymized_outline = anonymizer.anonymize_struct( outline, value_template=Template("{{ _${variable_name}_ }}") ) + anonymized_files = anonymizer.anonymize_struct( + files, value_template=Template("{{ _${variable_name}_ }}") + ) answer = { "role": anonymized_role, "outline": anonymized_outline, - "files": files, + "files": anonymized_files, "format": "plaintext", "generationId": self.validated_data["generationId"], } From 48a464444a7b94799239dd8c08a9957fed45c0d9 Mon Sep 17 00:00:00 2001 From: Michael Abashian Date: Wed, 22 Jan 2025 12:00:53 -0500 Subject: [PATCH 4/4] Update reverse function import --- ansible_ai_connect/ai/api/tests/test_role_generation_view.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ansible_ai_connect/ai/api/tests/test_role_generation_view.py b/ansible_ai_connect/ai/api/tests/test_role_generation_view.py index fe522bac1..603246e4e 100644 --- a/ansible_ai_connect/ai/api/tests/test_role_generation_view.py +++ b/ansible_ai_connect/ai/api/tests/test_role_generation_view.py @@ -21,7 +21,6 @@ from django.apps import apps from django.test import override_settings -from django.urls import reverse from ansible_ai_connect.ai.api.model_pipelines.config_pipelines import BaseConfig from ansible_ai_connect.ai.api.model_pipelines.dummy.pipelines import ROLE_FILES @@ -31,6 +30,7 @@ RoleGenerationResponse, ) from ansible_ai_connect.ai.api.model_pipelines.tests import mock_config +from ansible_ai_connect.ai.api.utils.version import api_version_reverse as reverse from ansible_ai_connect.healthcheck.backends import HealthCheckSummary from ansible_ai_connect.test_utils import ( WisdomAppsBackendMocking,