jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

fix unit test

Browse source
This commit is contained in:
ming luo 2024-06-21 20:46:54 -04:00 • committed by Gabriel Luiz Freitas Almeida
commit 17adc8b5a5
5 changed files with 162 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

6
poetry.lock generated
View file

@ -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"

8
src/backend/base/langflow/services/variable/base.py
View file

@ -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.

69
src/backend/base/langflow/services/variable/kubernetes_secrets.py
View file

@ -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

33
src/backend/base/langflow/services/variable/service.py
View file

@ -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

121
tests/unit/test_kubernetes_secrets.py
View file

@ -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