Merge pull request #203 from mantiumai/AlexN/implement-api-asset-ping

implement APIEndpointAsset ping functionality

chirps/asset/models.py

@@ -39,7 +39,9 @@ class BaseAsset(PolymorphicModel):
 
     def scan_is_active(self) -> bool:
         """Return True if the asset is currently being scanned."""
-        return self.scan_run_assets.filter(~Q(scan__status='Complete')).exists()
+        return self.scan_run_assets.filter(
+            ~Q(scan__status='Complete') & ~Q(scan__status='Failed') & ~Q(scan__status='Canceled')
+        ).exists()
 
     def search(self, query: str, max_results: int) -> list[SearchResult]:
         """Perform a query against the specified asset, returning the max_results number of matches."""

chirps/asset/providers/api_endpoint.py

@@ -45,8 +45,7 @@ def decrypted_api_key(self):
                 return 'Error: Decryption failed'
         return None
 
-    def fetch_api_data(self, query: str) -> dict:
-        """Fetch data from the API using the provided query."""
+    def _send_request(self, message: str) -> requests.Response:
         # Convert headers JSON string into a dictionary
         headers_dict = json.loads(self.headers) if self.headers else {}
 
@@ -58,14 +57,20 @@ def fetch_api_data(self, query: str) -> dict:
             headers['Authorization'] = f'Basic {self.api_key}'
 
         # Replace the %query% placeholder in the body
-        body = json.loads(json.dumps(self.body).replace('%query%', query))
+        body = json.loads(json.dumps(self.body).replace('%query%', message))
 
         # Send the request
         try:
             response = requests.post(self.url, headers=headers, json=body, timeout=self.timeout)
         except Timeout as exc:
             raise RequestException('Error: API request timed out') from exc
 
+        return response
+
+    def fetch_api_data(self, query: str) -> dict:
+        """Fetch data from the API using the provided query."""
+        response = self._send_request(query)
+
         # Check if the request was successful
         if response.status_code != 200:
             raise RequestException(f'Error: API request failed with status code {response.status_code}')
@@ -76,7 +81,21 @@ def fetch_api_data(self, query: str) -> dict:
 
     def test_connection(self) -> PingResult:
         """Ensure that the API Endpoint asset can be connected to."""
-        raise NotImplementedError('The test_connection method is not implemented for this asset.')
+        message = 'Test message'
+        response = self._send_request(message)
+
+        # Check if the request was successful and return a PingResult object
+        if response.status_code == 200:
+            return PingResult(success=True)
+
+        try:
+            response_data = response.json()
+            error_message = json.dumps(response_data, indent=2)
+        except ValueError:
+            error_message = response.text
+
+        error = f'Error: API request failed with status code {response.status_code} and response:\n{error_message}'
+        return PingResult(success=False, error=error)
 
     def displayable_attributes(self):
         """Display a subset of the model's attributes"""

chirps/asset/tests.py

@@ -4,6 +4,7 @@
 
 import fakeredis
 from asset.forms import APIEndpointAssetForm, PineconeAssetForm, RedisAssetForm
+from asset.models import PingResult
 from asset.providers.api_endpoint import APIEndpointAsset
 from asset.providers.mantium import MantiumAsset
 from asset.providers.redis import RedisAsset
@@ -296,3 +297,40 @@ def test_fetch_api_data(self):
             self.assertEqual(result['avatarFormat'], 'webm')
             self.assertEqual(result['secure'], 'true')
             self.assertEqual(result['message'], 'hello')
+
+    def test_test_connection(self):
+        """Test the test_connection method."""
+        with mock.patch('requests.post') as mock_post:
+            mock_post.return_value.status_code = 200
+
+            ping_result = self.api_endpoint_asset.test_connection()
+            self.assertIsInstance(ping_result, PingResult)
+            self.assertTrue(ping_result.success)
+            self.assertIsNone(ping_result.error)
+
+            expected_url = self.api_endpoint_asset.url
+            expected_headers = json.loads(self.api_endpoint_asset.headers)
+            expected_headers['Authorization'] = f'Bearer {self.api_endpoint_asset.api_key}'
+            expected_body = self.api_endpoint_asset.body.replace('%query%', 'Test message')
+
+            mock_post.assert_called_once_with(expected_url, headers=expected_headers, json=expected_body, timeout=30)
+
+    def test_test_connection_failure(self):
+        """Test the test_connection method when the API request fails."""
+        mock_response_data = {'error': 'API request failed'}
+
+        with mock.patch('requests.post') as mock_post:
+            mock_post.return_value.status_code = 400
+            mock_post.return_value.json.return_value = mock_response_data
+
+            ping_result = self.api_endpoint_asset.test_connection()
+            self.assertIsInstance(ping_result, PingResult)
+            self.assertFalse(ping_result.success)
+            self.assertIn(mock_response_data['error'], ping_result.error)
+
+            expected_url = self.api_endpoint_asset.url
+            expected_headers = json.loads(self.api_endpoint_asset.headers)
+            expected_headers['Authorization'] = f'Bearer {self.api_endpoint_asset.api_key}'
+            expected_body = self.api_endpoint_asset.body.replace('%query%', 'Test message')
+
+            mock_post.assert_called_once_with(expected_url, headers=expected_headers, json=expected_body, timeout=30)