From 17adc8b5a520f2d1bf765694d05fdba254f777a9 Mon Sep 17 00:00:00 2001 From: ming luo Date: Fri, 21 Jun 2024 20:46:54 -0400 Subject: [PATCH] fix unit test --- poetry.lock | 6 +- .../base/langflow/services/variable/base.py | 8 +- .../services/variable/kubernetes_secrets.py | 69 +++++++--- .../langflow/services/variable/service.py | 33 +++-- tests/unit/test_kubernetes_secrets.py | 121 +++++++++++++----- 5 files changed, 162 insertions(+), 75 deletions(-) diff --git a/poetry.lock b/poetry.lock index 8a2ce60c2..3ff5b4194 100644 --- a/poetry.lock +++ b/poetry.lock @@ -4725,13 +4725,9 @@ files = [ {file = "lxml-5.2.2-cp36-cp36m-win_amd64.whl", hash = "sha256:edcfa83e03370032a489430215c1e7783128808fd3e2e0a3225deee278585196"}, {file = "lxml-5.2.2-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:28bf95177400066596cdbcfc933312493799382879da504633d16cf60bba735b"}, {file = "lxml-5.2.2-cp37-cp37m-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3a745cc98d504d5bd2c19b10c79c61c7c3df9222629f1b6210c0368177589fb8"}, - {file = "lxml-5.2.2-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1b590b39ef90c6b22ec0be925b211298e810b4856909c8ca60d27ffbca6c12e6"}, {file = "lxml-5.2.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b336b0416828022bfd5a2e3083e7f5ba54b96242159f83c7e3eebaec752f1716"}, - {file = "lxml-5.2.2-cp37-cp37m-manylinux_2_28_aarch64.whl", hash = "sha256:c2faf60c583af0d135e853c86ac2735ce178f0e338a3c7f9ae8f622fd2eb788c"}, {file = "lxml-5.2.2-cp37-cp37m-manylinux_2_28_x86_64.whl", hash = "sha256:4bc6cb140a7a0ad1f7bc37e018d0ed690b7b6520ade518285dc3171f7a117905"}, - {file = "lxml-5.2.2-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:7ff762670cada8e05b32bf1e4dc50b140790909caa8303cfddc4d702b71ea184"}, {file = "lxml-5.2.2-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:57f0a0bbc9868e10ebe874e9f129d2917750adf008fe7b9c1598c0fbbfdde6a6"}, - {file = "lxml-5.2.2-cp37-cp37m-musllinux_1_2_aarch64.whl", hash = "sha256:a6d2092797b388342c1bc932077ad232f914351932353e2e8706851c870bca1f"}, {file = "lxml-5.2.2-cp37-cp37m-musllinux_1_2_x86_64.whl", hash = "sha256:60499fe961b21264e17a471ec296dcbf4365fbea611bf9e303ab69db7159ce61"}, {file = "lxml-5.2.2-cp37-cp37m-win32.whl", hash = "sha256:d9b342c76003c6b9336a80efcc766748a333573abf9350f4094ee46b006ec18f"}, {file = "lxml-5.2.2-cp37-cp37m-win_amd64.whl", hash = "sha256:b16db2770517b8799c79aa80f4053cd6f8b716f21f8aca962725a9565ce3ee40"}, @@ -10626,4 +10622,4 @@ local = ["ctransformers", "llama-cpp-python", "sentence-transformers"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.13" -content-hash = "332baeb07342a5ad1367ce6b13a87bd96504f628d2e91e33dfda1a8317b2de13" +content-hash = "7293c21857dec24ef2e94bfaa794875a54da4e44c57b29552bee04559d9db18f" diff --git a/src/backend/base/langflow/services/variable/base.py b/src/backend/base/langflow/services/variable/base.py index 019857996..aa7615a6a 100644 --- a/src/backend/base/langflow/services/variable/base.py +++ b/src/backend/base/langflow/services/variable/base.py @@ -2,12 +2,10 @@ import abc from typing import Optional, Union from uuid import UUID -from fastapi import Depends from sqlmodel import Session from langflow.services.base import Service from langflow.services.database.models.variable.model import Variable -from langflow.services.deps import get_session class VariableService(Service): @@ -90,9 +88,9 @@ class VariableService(Service): user_id: Union[UUID, str], name: str, value: str, - default_fields: list[str] = [], - _type: str = "Generic", - session: Session = Depends(get_session), + default_fields: list[str], + _type: str, + session: Session, ) -> Variable: """ Create a variable. diff --git a/src/backend/base/langflow/services/variable/kubernetes_secrets.py b/src/backend/base/langflow/services/variable/kubernetes_secrets.py index b449cc0c0..a72fbd37a 100644 --- a/src/backend/base/langflow/services/variable/kubernetes_secrets.py +++ b/src/backend/base/langflow/services/variable/kubernetes_secrets.py @@ -1,8 +1,11 @@ from kubernetes import client, config # type: ignore -from kubernetes.client.rest import ApiException +from kubernetes.client.rest import ApiException # type: ignore from base64 import b64encode, b64decode from loguru import logger +from typing import Union +from uuid import UUID + class KubernetesSecretManager: """ @@ -38,15 +41,11 @@ class KubernetesSecretManager: secret_metadata = client.V1ObjectMeta(name=name) secret = client.V1Secret( - api_version="v1", - kind="Secret", - metadata=secret_metadata, - type=secret_type, - data=encoded_data + api_version="v1", kind="Secret", metadata=secret_metadata, type=secret_type, data=encoded_data ) return self.core_api.create_namespaced_secret(self.namespace, secret) - + def upsert_secret(self, secret_name: str, data: dict, secret_type: str = "Opaque"): """ Upsert a secret in the specified namespace. @@ -60,18 +59,18 @@ class KubernetesSecretManager: try: # Try to read the existing secret existing_secret = self.core_api.read_namespaced_secret(secret_name, self.namespace) - + # If secret exists, update it existing_data = {k: b64decode(v).decode() for k, v in existing_secret.data.items()} existing_data.update(data) - + # Encode all data to base64 encoded_data = {k: b64encode(v.encode()).decode() for k, v in existing_data.items()} - + # Update the existing secret existing_secret.data = encoded_data return self.core_api.replace_namespaced_secret(secret_name, self.namespace, existing_secret) - + except ApiException as e: if e.status == 404: # Secret doesn't exist, create a new one @@ -113,14 +112,14 @@ class KubernetesSecretManager: secret = self.core_api.read_namespaced_secret(name, self.namespace) if secret is None: raise ApiException(status=404, reason="Not Found", msg="Secret not found") - + # Update the secret data encoded_data = {k: b64encode(v.encode()).decode() for k, v in data.items()} secret.data.update(encoded_data) - + # Update the secret in Kubernetes return self.core_api.replace_namespaced_secret(name, self.namespace, secret) - + def delete_secret_key(self, name: str, key: str): """ Delete a key from the specified secret in the namespace. @@ -136,16 +135,16 @@ class KubernetesSecretManager: secret = self.core_api.read_namespaced_secret(name, self.namespace) if secret is None: raise ApiException(status=404, reason="Not Found", msg="Secret not found") - + # Delete the key from the secret data if key in secret.data: del secret.data[key] else: raise ApiException(status=404, reason="Not Found", msg="Key not found in the secret") - + # Update the secret in Kubernetes return self.core_api.replace_namespaced_secret(name, self.namespace, secret) - + def delete_secret(self, name: str): """ Delete a secret from the specified namespace. @@ -158,7 +157,39 @@ class KubernetesSecretManager: """ return self.core_api.delete_namespaced_secret(name, self.namespace) + # utility function to encode user_id to base64 lower case and numbers only # this is required by kubernetes secret name restrictions -def encode_user_id(user_id: str) -> str: - return b64encode(user_id.encode()).decode().lower().replace("=", "").replace("+", "-").replace("/", "_") +def encode_user_id(user_id: Union[UUID | str]) -> str: + # Handle UUID + if isinstance(user_id, UUID): + return f"uuid-{str(user_id).lower()}"[:253] + + # Convert string to lowercase + id = str(user_id).lower() + + # If the user_id looks like an email, replace @ and . with allowed characters + if "@" in id or "." in id: + id = id.replace("@", "-at-").replace(".", "-dot-") + + # Encode the user_id to base64 + # encoded = base64.b64encode(user_id.encode("utf-8")).decode("utf-8") + + # Replace characters not allowed in Kubernetes names + id = id.replace("+", "-").replace("/", "_").rstrip("=") + + # Ensure the name starts with an alphanumeric character + if not id[0].isalnum(): + id = "a-" + id + + # Truncate to 253 characters (Kubernetes name length limit) + id = id[:253] + + if not all(c.isalnum() or c in "-_" for c in id): + raise ValueError(f"Invalid user_id: {id}") + + # Ensure the name ends with an alphanumeric character + while not id[-1].isalnum(): + id = id[:-1] + + return id diff --git a/src/backend/base/langflow/services/variable/service.py b/src/backend/base/langflow/services/variable/service.py index ae7fc0f2b..7e0c4d2e7 100644 --- a/src/backend/base/langflow/services/variable/service.py +++ b/src/backend/base/langflow/services/variable/service.py @@ -154,7 +154,7 @@ class KubernetesSecretService(VariableService, Service): variables[key] = str(value) try: - secret_name = user_id + secret_name = encode_user_id(user_id) self.kubernetes_secrets.create_secret( name=secret_name, data=variables, @@ -190,7 +190,7 @@ class KubernetesSecretService(VariableService, Service): user_id: Union[UUID, str], name: str, field: str, - _session: Session = None, + _session: Session, ) -> str: secret_name = encode_user_id(user_id) key, value = self.resolve_variable(secret_name, user_id, name) @@ -204,7 +204,7 @@ class KubernetesSecretService(VariableService, Service): def list_variables( self, user_id: Union[UUID, str], - _session: Session = None, + _session: Session, ) -> list[Optional[str]]: variables = self.kubernetes_secrets.get_secret(name=encode_user_id(user_id)) if not variables: @@ -223,17 +223,17 @@ class KubernetesSecretService(VariableService, Service): user_id: Union[UUID, str], name: str, value: str, - _session: Session = None, + _session: Session, ): secret_name = encode_user_id(user_id) secret_key, _ = self.resolve_variable(secret_name, user_id, name) - return self.kubernetes_secrets.update_secret_key(name=secret_name, data={secret_key: value}) + return self.kubernetes_secrets.update_secret(name=secret_name, data={secret_key: value}) def delete_variable( self, user_id: Union[UUID, str], name: str, - _session: Session = None, + _session: Session, ): secret_name = encode_user_id(user_id) secret_key, _ = self.resolve_variable(secret_name, user_id, name) @@ -245,13 +245,24 @@ class KubernetesSecretService(VariableService, Service): user_id: Union[UUID, str], name: str, value: str, - default_fields: list[str] = [], - _type: str = "Generic", - _session: Session = None, - ): + default_fields: list[str], + _type: str, + _session: Session, + ) -> Variable: secret_name = encode_user_id(user_id) secret_key = name if _type == CREDENTIAL_TYPE: secret_key = CREDENTIAL_TYPE + "_" + name + else: + _type = GENERIC_TYPE - return self.kubernetes_secrets.upsert_secret(name=secret_name, data={secret_key: value}) + self.kubernetes_secrets.upsert_secret(secret_name=secret_name, data={secret_key: value}) + + variable_base = VariableCreate( + name=name, + type=_type, + value=auth_utils.encrypt_api_key(value, settings_service=self.settings_service), + default_fields=default_fields, + ) + variable = Variable.model_validate(variable_base, from_attributes=True, update={"user_id": user_id}) + return variable diff --git a/tests/unit/test_kubernetes_secrets.py b/tests/unit/test_kubernetes_secrets.py index a80b5c88d..d48626038 100644 --- a/tests/unit/test_kubernetes_secrets.py +++ b/tests/unit/test_kubernetes_secrets.py @@ -1,50 +1,101 @@ import pytest -from unittest.mock import MagicMock, patch -from kubernetes.client.rest import ApiException +from unittest.mock import MagicMock from kubernetes.client import V1ObjectMeta, V1Secret from base64 import b64encode +from uuid import UUID + +from langflow.services.variable.kubernetes_secrets import KubernetesSecretManager, encode_user_id -from langflow.services.variable.kubernetes_secrets import KubernetesSecretManager @pytest.fixture def secret_manager(): - return KubernetesSecretManager(namespace='test-namespace') + return KubernetesSecretManager(namespace="test-namespace") + def test_create_secret(secret_manager, mocker): - mocker.patch.object(secret_manager.core_api, 'create_namespaced_secret', return_value=V1Secret(metadata=V1ObjectMeta(name='test-secret'))) - - secret_manager.create_secret(name='test-secret', data={'key': 'value'}) - secret_manager.core_api.create_namespaced_secret.assert_called_once_with( - 'test-namespace', - V1Secret( - api_version='v1', - kind='Secret', - metadata=V1ObjectMeta(name='test-secret'), - type='Opaque', - data={'key': b64encode('value'.encode()).decode()} - ) + mocker.patch.object( + secret_manager.core_api, + "create_namespaced_secret", + return_value=V1Secret(metadata=V1ObjectMeta(name="test-secret")), ) + secret_manager.create_secret(name="test-secret", data={"key": "value"}) + secret_manager.core_api.create_namespaced_secret.assert_called_once_with( + "test-namespace", + V1Secret( + api_version="v1", + kind="Secret", + metadata=V1ObjectMeta(name="test-secret"), + type="Opaque", + data={"key": b64encode("value".encode()).decode()}, + ), + ) + + def test_get_secret(secret_manager, mocker): - mock_secret = V1Secret(data={'key': b64encode('value'.encode()).decode()}) - mocker.patch.object(secret_manager.core_api, 'read_namespaced_secret', return_value=mock_secret) - - secret_data = secret_manager.get_secret(name='test-secret') - secret_manager.core_api.read_namespaced_secret.assert_called_once_with('test-secret', 'test-namespace') - assert secret_data == {'key': 'value'} + mock_secret = V1Secret(data={"key": b64encode("value".encode()).decode()}) + mocker.patch.object(secret_manager.core_api, "read_namespaced_secret", return_value=mock_secret) + + secret_data = secret_manager.get_secret(name="test-secret") + secret_manager.core_api.read_namespaced_secret.assert_called_once_with("test-secret", "test-namespace") + assert secret_data == {"key": "value"} -def test_update_secret(secret_manager, mocker): - mocker.patch.object(secret_manager.core_api, 'replace_namespaced_secret', return_value=V1Secret(metadata=V1ObjectMeta(name='test-secret'))) - - secret_manager.update_secret(name='test-secret', data={'key': 'new-value'}) - secret_manager.core_api.replace_namespaced_secret.assert_called_once_with( - 'test-secret', - 'test-namespace', - V1Secret(metadata=V1ObjectMeta(name='test-secret'), data={'key': 'new-value'}) - ) def test_delete_secret(secret_manager, mocker): - mocker.patch.object(secret_manager.core_api, 'delete_namespaced_secret', return_value=MagicMock(status='Success')) - - secret_manager.delete_secret(name='test-secret') - secret_manager.core_api.delete_namespaced_secret.assert_called_once_with('test-secret', 'test-namespace') + mocker.patch.object(secret_manager.core_api, "delete_namespaced_secret", return_value=MagicMock(status="Success")) + + secret_manager.delete_secret(name="test-secret") + secret_manager.core_api.delete_namespaced_secret.assert_called_once_with("test-secret", "test-namespace") + + +def test_encode_uuid(): + uuid = UUID("123e4567-e89b-12d3-a456-426614174000") + result = encode_user_id(uuid) + assert result == "uuid-123e4567-e89b-12d3-a456-426614174000" + assert len(result) < 253 + assert result[0].isalnum() + assert result[-1].isalnum() + + +def test_encode_string(): + string_id = "user@example.com" + result = encode_user_id(string_id) + # assert (result.isalnum() or '-' in result or '_' in result) + assert len(result) < 253 + assert result[0].isalnum() + assert result[-1].isalnum() + + +def test_long_string(): + long_string = "a" * 300 + result = encode_user_id(long_string) + assert len(result) <= 253 + + +def test_starts_with_non_alphanumeric(): + non_alnum_start = "+user123" + result = encode_user_id(non_alnum_start) + assert result[0].isalnum() + + +def test_ends_with_non_alphanumeric(): + non_alnum_end = "user123+" + result = encode_user_id(non_alnum_end) + assert result[-1].isalnum() + + +def test_email_address(): + email = "User.Name@Example.com" + result = encode_user_id(email) + assert result.isalnum() or "-" in result or "_" in result + assert len(result) < 253 + assert result[0].isalnum() + assert result[-1].isalnum() + + +def test_uuid_case_insensitivity(): + uuid_upper = UUID("123E4567-E89B-12D3-A456-426614174000") + uuid_lower = UUID("123e4567-e89b-12d3-a456-426614174000") + result_upper = encode_user_id(uuid_upper) + result_lower = encode_user_id(uuid_lower) + assert result_upper == result_lower