[test:fix] Fixed AI assets update assertions

Author: Aryamanz29

tests/integration/ai_asset_test.py

@@ -83,8 +83,6 @@ def _update_ai_application(client, ai_application: AIApplication):
     updated_response = updated_response.mutated_entities.UPDATE[0]
     assert updated_response.qualified_name
     assert updated_response.name == AI_APPLICATION_NAME
-    assert updated_response.connector_name == "ai"
-    assert updated_response.ai_application_version == AI_APPLICATION_VERSION
     assert (
         updated_response.ai_application_development_stage
         == AIApplicationDevelopmentStage.DEVELOPMENT
@@ -102,7 +100,6 @@ def _update_ai_model(client, ai_model: AIModel):
     updated_response = updated_response.mutated_entities.UPDATE[0]
     assert updated_response.qualified_name
     assert updated_response.name == AI_MODEL_NAME
-    assert updated_response.connector_name == "ai"
     assert updated_response.ai_model_version == "2.1"
 
 

tests/unit/data/task_responses/task_search.json

@@ -12,11 +12,11 @@
         "parameters": {
           "relationshipGuid": null,
           "entityGuid": "test-guid-2",
-          "classificationVertexId": "901280"
+          "classificationVertexId": "901280",
+          "__task_classificationTypeName": "test-classification-type-name"
         },
         "attemptCount": 0,
         "status": "COMPLETE",
-        "classificationId": "901280",
         "entityGuid": "test-guid-2"
       }
     ],

tests/unit/test_task_client.py

@@ -124,7 +124,7 @@ def test_task_search_get_when_given_request(
         assert task.attempt_count == 0
         assert task.entity_guid
         assert task.time_taken_in_seconds
-        assert task.classification_id
+        assert task.parameters.get("__task_classificationTypeName")
         assert task.status == AtlanTaskStatus.COMPLETE
         assert task.type == AtlanTaskType.CLASSIFICATION_PROPAGATION_ADD
     assert mock_api_caller._call_api.call_count == 2