aixtools 0.1.6__tar.gz → 0.1.7__tar.gz

{aixtools-0.1.6 → aixtools-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.1.6
+Version: 0.1.7
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown

{aixtools-0.1.6 → aixtools-0.1.7}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.6'
-__version_tuple__ = version_tuple = (0, 1, 6)
+__version__ = version = '0.1.7'
+__version_tuple__ = version_tuple = (0, 1, 7)
 
-__commit_id__ = commit_id = 'g6abaa9c74'
+__commit_id__ = commit_id = 'g78a5dc717'

aixtools-0.1.7/aixtools/vault/vault.py

@@ -0,0 +1,100 @@
+# ruff: noqa: PLR0913
+"""
+Provides a Vault client for storing and retrieving user service api keys.
+"""
+
+import logging
+from typing import Dict, Optional
+
+import hvac
+from hvac.exceptions import InvalidPath
+
+from aixtools.utils import config
+
+logger = logging.getLogger(__name__)
+
+
+class VaultAuthError(Exception):
+    """Exception raised for vault authentication errors."""
+
+
+class VaultClient:
+    """Vault client for storing and retrieving user service api keys."""
+
+    def __init__(self):
+        self.client = hvac.Client(url=config.VAULT_ADDRESS, token=config.VAULT_TOKEN)
+
+        if not self.client.is_authenticated():
+            raise VaultAuthError("Vault client authentication failed. Check vault_token.")
+
+    def _get_secret_path(self, user_id: str, service_name: str) -> str:
+        """Generate the vault secret path for a user and service."""
+        return f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
+
+    def store_user_service_api_key(self, *, user_id: str, service_name: str, user_api_key: str):
+        """
+        Store user's service api key in the Vault at the specified vault mount
+        point, where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+
+        This is a convenience method for storing a single API key.
+        For storing multiple secrets, use store_user_service_secret().
+        """
+        secret_dict = {"user-api-key": user_api_key}
+        self.store_user_service_secret(user_id=user_id, service_name=service_name, secret_data=secret_dict)
+
+    def read_user_service_api_key(self, *, user_id: str, service_name: str) -> Optional[str]:
+        """
+        Read user's service api key in from vault at the specified mount point,
+        where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+
+        This is a convenience method for reading a single API key.
+        For reading multiple secrets, use read_user_service_secret().
+        """
+        secret_data = self.read_user_service_secret(user_id=user_id, service_name=service_name)
+        if secret_data is None:
+            return None
+        return secret_data.get("user-api-key")
+
+    def store_user_service_secret(self, *, user_id: str, service_name: str, secret_data: Dict[str, str]):
+        """
+        Store complete user service secret with multiple key-value pairs in the Vault
+        at the specified vault mount point, where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+        """
+        secret_path = None
+        try:
+            secret_path = self._get_secret_path(user_id, service_name)
+            logger.info("Writing complete secret to path %s", secret_path)
+            self.client.secrets.kv.v2.create_or_update_secret(
+                secret_path, secret=secret_data, mount_point=config.VAULT_MOUNT_POINT
+            )
+
+            logger.info("Complete secret written to path %s", secret_path)
+        except Exception as e:
+            logger.error("Failed to write complete secret to path %s: %s", secret_path, str(e))
+            raise VaultAuthError(e) from e
+
+    def read_user_service_secret(self, *, user_id: str, service_name: str) -> Optional[Dict[str, str]]:
+        """
+        Read complete user service secret from vault at the specified mount point,
+        where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+        Returns all key-value pairs in the secret or None if the secret doesn't exist.
+        """
+        secret_path = None
+
+        try:
+            secret_path = self._get_secret_path(user_id, service_name)
+            logger.info("Reading complete secret from path %s", secret_path)
+            response = self.client.secrets.kv.v2.read_secret_version(
+                secret_path, mount_point=config.VAULT_MOUNT_POINT, raise_on_deleted_version=True
+            )
+            secret_data = response["data"]["data"]
+            logger.info("Complete secret read from path %s", secret_path)
+            return secret_data
+        except InvalidPath:
+            # Secret path does not exist
+            logger.warning("Secret path does not exist %s", secret_path)
+            return None
+
+        except Exception as e:
+            logger.error("Failed to read complete secret from path %s: %s", secret_path, str(e))
+            raise VaultAuthError(e) from e

{aixtools-0.1.6 → aixtools-0.1.7}/aixtools.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.1.6
+Version: 0.1.7
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown

aixtools-0.1.7/tests/unit/vault/test_vault.py

@@ -0,0 +1,171 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from hvac.exceptions import InvalidPath
+
+from aixtools.vault.vault import VaultAuthError, VaultClient
+
+
+@pytest.fixture
+def patched_vault_client():
+    with patch("aixtools.vault.vault.hvac.Client") as mock_hvac_client_cls:
+        fake_hvac_client = MagicMock()
+        fake_hvac_client.is_authenticated.return_value = True
+        mock_hvac_client_cls.return_value = fake_hvac_client
+
+        client = VaultClient()
+        return client
+
+
+@pytest.fixture
+def valid_params():
+    return {
+        "vault_mount_point": "secret",
+        "path_prefix": "path",
+        "env": "dev",
+        "user_id": "test-user",
+        "service_name": "test-service",
+        "user_api_key": "test-api-key",
+    }
+
+
+def test_store_user_service_api_key_success(patched_vault_client, valid_params):
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.return_value = {}
+
+    patched_vault_client.store_user_service_api_key(
+        user_id=valid_params["user_id"],
+        service_name=valid_params["service_name"],
+        user_api_key=valid_params["user_api_key"],
+    )
+
+    secret_path = (
+        f"{valid_params['path_prefix']}/{valid_params['env']}/{valid_params['user_id']}/{valid_params['service_name']}"
+    )
+
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.assert_called_once_with(
+        secret_path,
+        secret={"user-api-key": valid_params["user_api_key"]},
+        mount_point=valid_params["vault_mount_point"],
+    )
+
+
+def test_store_user_service_api_key_invalid_path(patched_vault_client, valid_params):
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.side_effect = Exception("Invalid Path")
+
+    with pytest.raises(VaultAuthError, match="Invalid Path"):
+        patched_vault_client.store_user_service_api_key(
+            user_id=valid_params["user_id"],
+            service_name=valid_params["service_name"],
+            user_api_key=valid_params["user_api_key"],
+        )
+        patched_vault_client.client.assert_not_called()
+
+
+def test_read_user_service_api_key_success(patched_vault_client):
+    """Test successful read of user service API key."""
+    mock_response = {"data": {"data": {"user-api-key": "test-api-key"}}}
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.return_value = mock_response
+
+    result = patched_vault_client.read_user_service_api_key(user_id="test-user", service_name="test-service")
+
+    assert result == "test-api-key"
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
+    )
+
+
+def test_read_user_service_api_key_secret_not_found(patched_vault_client):
+    """Test read_user_service_api_key when the secret path does not exist."""
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = InvalidPath
+
+    result = patched_vault_client.read_user_service_api_key(user_id="test-user", service_name="test-service")
+
+    assert result is None
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
+    )
+
+
+def test_read_user_service_api_key_unexpected_error(patched_vault_client):
+    """Test read_user_service_api_key when an unexpected exception occurs."""
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = Exception("Unexpected error")
+
+    with pytest.raises(VaultAuthError, match="Unexpected error"):
+        patched_vault_client.read_user_service_api_key(user_id="test-user", service_name="test-service")
+
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
+    )
+
+
+def test_store_user_service_secret_success(patched_vault_client, valid_params):
+    """Test successful storage of complete user service secret."""
+    secret_data = {"api_key": "test-api-key", "token": "test-token", "endpoint": "https://api.example.com"}
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.return_value = {}
+
+    patched_vault_client.store_user_service_secret(
+        user_id=valid_params["user_id"],
+        service_name=valid_params["service_name"],
+        secret_data=secret_data,
+    )
+
+    secret_path = (
+        f"{valid_params['path_prefix']}/{valid_params['env']}/{valid_params['user_id']}/{valid_params['service_name']}"
+    )
+
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.assert_called_once_with(
+        secret_path,
+        secret=secret_data,
+        mount_point=valid_params["vault_mount_point"],
+    )
+
+
+def test_store_user_service_secret_error(patched_vault_client, valid_params):
+    """Test store_user_service_secret when an error occurs."""
+    secret_data = {"api_key": "test-api-key", "token": "test-token"}
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.side_effect = Exception("Storage error")
+
+    with pytest.raises(VaultAuthError, match="Storage error"):
+        patched_vault_client.store_user_service_secret(
+            user_id=valid_params["user_id"],
+            service_name=valid_params["service_name"],
+            secret_data=secret_data,
+        )
+
+
+def test_read_user_service_secret_success(patched_vault_client):
+    """Test successful read of complete user service secret."""
+    secret_data = {"api_key": "test-api-key", "token": "test-token", "endpoint": "https://api.example.com"}
+    mock_response = {"data": {"data": secret_data}}
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.return_value = mock_response
+
+    result = patched_vault_client.read_user_service_secret(user_id="test-user", service_name="test-service")
+
+    assert result == secret_data
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
+    )
+
+
+def test_read_user_service_secret_not_found(patched_vault_client):
+    """Test read_user_service_secret when the secret path does not exist."""
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = InvalidPath
+
+    result = patched_vault_client.read_user_service_secret(user_id="test-user", service_name="test-service")
+
+    assert result is None
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
+    )
+
+
+def test_read_user_service_secret_error(patched_vault_client):
+    """Test read_user_service_secret when an unexpected exception occurs."""
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = Exception("Read error")
+
+    with pytest.raises(VaultAuthError, match="Read error"):
+        patched_vault_client.read_user_service_secret(user_id="test-user", service_name="test-service")
+
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
+    )

aixtools-0.1.6/aixtools/vault/vault.py

@@ -1,73 +0,0 @@
-# ruff: noqa: PLR0913
-"""
-Provides a Vault client for storing and retrieving user service api keys.
-"""
-
-import logging
-from typing import Optional
-
-import hvac
-from hvac.exceptions import InvalidPath
-
-from aixtools.utils import config
-
-logger = logging.getLogger(__name__)
-
-
-class VaultAuthError(Exception):
-    """Exception raised for vault authentication errors."""
-
-
-class VaultClient:
-    """Vault client for storing and retrieving user service api keys."""
-
-    def __init__(self):
-        self.client = hvac.Client(url=config.VAULT_ADDRESS, token=config.VAULT_TOKEN)
-
-        if not self.client.is_authenticated():
-            raise VaultAuthError("Vault client authentication failed. Check vault_token.")
-
-    def store_user_service_api_key(self, *, user_id: str, service_name: str, user_api_key: str):
-        """
-        Store user's service api key in the Vault at the specified vault mount
-        point, where the path is <path_prefix>/<env>/<user_id>/<service_name>.
-        """
-        secret_path = None
-        try:
-            secret_path = f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
-            print("secret_path", secret_path)
-            secret_dict = {"user-api-key": user_api_key}
-            self.client.secrets.kv.v2.create_or_update_secret(
-                secret_path, secret=secret_dict, mount_point=config.VAULT_MOUNT_POINT
-            )
-
-            logger.info("Secret written to path %s", secret_path)
-        except Exception as e:
-            logger.error("Failed to write secret to path %s: %s", secret_path, str(e))
-            raise VaultAuthError(e) from e
-
-    def read_user_service_api_key(self, *, user_id: str, service_name) -> Optional[str]:
-        """
-        Read user's service api key in from vault at the specified mount point,
-        where the path is <path_prefix>/<env>/<user_id>/<service_name>.
-        """
-        secret_path = None
-
-        try:
-            secret_path = f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
-            logger.info("Reading secret from path %s", secret_path)
-            response = self.client.secrets.kv.v2.read_secret_version(
-                secret_path, mount_point=config.VAULT_MOUNT_POINT, raise_on_deleted_version=True
-            )
-            secret_data = response["data"]["data"]
-            user_api_key = secret_data["user-api-key"]
-            logger.info("Secret read from path %s ", secret_path)
-            return user_api_key
-        except InvalidPath:
-            # Secret path does not exist
-            logger.warning("Secret path does not exist %s ", secret_path)
-            return None
-
-        except Exception as e:
-            logger.error("Failed to read secret from path %s: %s", secret_path, str(e))
-            raise VaultAuthError(e) from e

aixtools-0.1.6/tests/unit/vault/test_vault.py

@@ -1,114 +0,0 @@
-from unittest.mock import MagicMock, patch
-
-import pytest
-from hvac.exceptions import InvalidPath
-
-from aixtools.vault.vault import VaultClient, VaultAuthError
-
-
-@pytest.fixture
-def patched_vault_client():
-    with patch("aixtools.vault.vault.hvac.Client") as mock_hvac_client_cls:
-        fake_hvac_client = MagicMock()
-        fake_hvac_client.is_authenticated.return_value = True
-        mock_hvac_client_cls.return_value = fake_hvac_client
-
-        client = VaultClient()
-        return client
-
-
-@pytest.fixture
-def valid_params():
-    return {
-        "vault_mount_point": "secret",
-        "path_prefix": "path",
-        "env": "dev",
-        "user_id": "test-user",
-        "service_name": "test-service",
-        "user_api_key": "test-api-key"
-    }
-
-
-def test_store_user_service_api_key_success(patched_vault_client, valid_params):
-    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.return_value = {}
-
-    patched_vault_client.store_user_service_api_key(user_id=valid_params["user_id"],
-                                                    service_name=valid_params['service_name'],
-                                                    user_api_key=valid_params['user_api_key'])
-
-    secret_path = f"{valid_params['path_prefix']}/{valid_params['env']}/{valid_params['user_id']}/{valid_params['service_name']}"
-
-    print("expected secret path", secret_path)
-    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.assert_called_once_with(
-        secret_path,
-        secret={"user-api-key": valid_params["user_api_key"]},
-        mount_point=valid_params["vault_mount_point"]
-    )
-
-
-def test_store_user_service_api_key_invalid_path(patched_vault_client, valid_params):
-    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.side_effect = Exception("Invalid Path")
-
-    with pytest.raises(VaultAuthError, match="Invalid Path"):
-        patched_vault_client.store_user_service_api_key(user_id=valid_params["user_id"],
-                                                    service_name=valid_params['service_name'],
-                                                    user_api_key=valid_params['user_api_key'])
-        patched_vault_client.client.assert_not_called()
-
-
-def test_read_user_service_api_key_success(patched_vault_client):
-    """Test successful read of user service API key."""
-    mock_response = {
-        "data": {
-            "data": {
-                "user-api-key": "test-api-key"
-            }
-        }
-    }
-    patched_vault_client.client.secrets.kv.v2.read_secret_version.return_value = mock_response
-
-    result = patched_vault_client.read_user_service_api_key(
-        user_id="test-user",
-        service_name="test-service"
-    )
-
-    assert result == "test-api-key"
-    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
-        "path/dev/test-user/test-service",
-        mount_point='secret',
-        raise_on_deleted_version=True
-    )
-
-
-def test_read_user_service_api_key_secret_not_found(patched_vault_client):
-    """Test read_user_service_api_key when the secret path does not exist."""
-    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = InvalidPath
-
-    result = patched_vault_client.read_user_service_api_key(
-        user_id="test-user",
-        service_name="test-service"
-    )
-
-    assert result is None
-    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
-        "path/dev/test-user/test-service",
-        mount_point="secret",
-        raise_on_deleted_version=True
-    )
-
-
-def test_read_user_service_api_key_unexpected_error(patched_vault_client):
-    """Test read_user_service_api_key when an unexpected exception occurs."""
-    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = Exception("Unexpected error")
-
-    with pytest.raises(VaultAuthError, match="Unexpected error"):
-        patched_vault_client.read_user_service_api_key(
-            user_id="test-user",
-            service_name="test-service"
-        )
-
-    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
-        "path/dev/test-user/test-service",
-        mount_point="secret",
-        raise_on_deleted_version=True
-    )