lionagi 0.0.208__py3-none-any.whl → 0.0.210__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- lionagi/__init__.py +4 -6
- lionagi/api_service/base_endpoint.py +65 -0
- lionagi/api_service/base_rate_limiter.py +121 -0
- lionagi/api_service/base_service.py +146 -0
- lionagi/api_service/chat_completion.py +6 -0
- lionagi/api_service/embeddings.py +6 -0
- lionagi/api_service/payload_package.py +47 -0
- lionagi/api_service/status_tracker.py +29 -0
- lionagi/core/__init__.py +5 -9
- lionagi/core/branch.py +1191 -0
- lionagi/core/flow.py +423 -0
- lionagi/core/{instruction_set/instruction_set.py → instruction_set.py} +3 -3
- lionagi/core/session.py +872 -0
- lionagi/schema/__init__.py +5 -8
- lionagi/schema/base_schema.py +821 -0
- lionagi/{_services → services}/base_service.py +4 -4
- lionagi/{_services → services}/oai.py +4 -4
- lionagi/structures/graph.py +1 -1
- lionagi/structures/relationship.py +1 -1
- lionagi/structures/structure.py +1 -1
- lionagi/tools/tool_manager.py +0 -163
- lionagi/tools/tool_util.py +2 -1
- lionagi/utils/__init__.py +7 -14
- lionagi/utils/api_util.py +63 -2
- lionagi/utils/core_utils.py +338 -0
- lionagi/utils/sys_util.py +3 -3
- lionagi/version.py +1 -1
- {lionagi-0.0.208.dist-info → lionagi-0.0.210.dist-info}/METADATA +28 -29
- lionagi-0.0.210.dist-info/RECORD +56 -0
- lionagi/_services/anthropic.py +0 -79
- lionagi/_services/anyscale.py +0 -0
- lionagi/_services/azure.py +0 -1
- lionagi/_services/bedrock.py +0 -0
- lionagi/_services/everlyai.py +0 -0
- lionagi/_services/gemini.py +0 -0
- lionagi/_services/gpt4all.py +0 -0
- lionagi/_services/huggingface.py +0 -0
- lionagi/_services/litellm.py +0 -33
- lionagi/_services/localai.py +0 -0
- lionagi/_services/openllm.py +0 -0
- lionagi/_services/openrouter.py +0 -44
- lionagi/_services/perplexity.py +0 -0
- lionagi/_services/predibase.py +0 -0
- lionagi/_services/rungpt.py +0 -0
- lionagi/_services/vllm.py +0 -0
- lionagi/_services/xinference.py +0 -0
- lionagi/agents/planner.py +0 -1
- lionagi/agents/prompter.py +0 -1
- lionagi/agents/scorer.py +0 -1
- lionagi/agents/summarizer.py +0 -1
- lionagi/agents/validator.py +0 -1
- lionagi/bridge/__init__.py +0 -22
- lionagi/bridge/langchain.py +0 -195
- lionagi/bridge/llama_index.py +0 -266
- lionagi/core/branch/__init__.py +0 -0
- lionagi/core/branch/branch.py +0 -841
- lionagi/core/branch/cluster.py +0 -1
- lionagi/core/branch/conversation.py +0 -787
- lionagi/core/core_util.py +0 -0
- lionagi/core/flow/__init__.py +0 -0
- lionagi/core/flow/flow.py +0 -19
- lionagi/core/flow/flow_util.py +0 -62
- lionagi/core/instruction_set/__init__.py +0 -0
- lionagi/core/messages/__init__.py +0 -0
- lionagi/core/sessions/__init__.py +0 -0
- lionagi/core/sessions/session.py +0 -504
- lionagi/datastores/__init__.py +0 -1
- lionagi/datastores/chroma.py +0 -1
- lionagi/datastores/deeplake.py +0 -1
- lionagi/datastores/elasticsearch.py +0 -1
- lionagi/datastores/lantern.py +0 -1
- lionagi/datastores/pinecone.py +0 -1
- lionagi/datastores/postgres.py +0 -1
- lionagi/datastores/qdrant.py +0 -1
- lionagi/loaders/__init__.py +0 -18
- lionagi/loaders/chunker.py +0 -166
- lionagi/loaders/load_util.py +0 -240
- lionagi/loaders/reader.py +0 -122
- lionagi/models/__init__.py +0 -0
- lionagi/models/base_model.py +0 -0
- lionagi/models/imodel.py +0 -53
- lionagi/schema/async_queue.py +0 -158
- lionagi/schema/base_condition.py +0 -1
- lionagi/schema/base_node.py +0 -422
- lionagi/schema/base_tool.py +0 -44
- lionagi/schema/data_logger.py +0 -126
- lionagi/schema/data_node.py +0 -88
- lionagi/schema/status_tracker.py +0 -37
- lionagi/tests/test_utils/test_encrypt_util.py +0 -323
- lionagi/utils/encrypt_util.py +0 -283
- lionagi/utils/url_util.py +0 -55
- lionagi-0.0.208.dist-info/RECORD +0 -106
- lionagi/{agents → api_service}/__init__.py +0 -0
- lionagi/core/{branch/branch_manager.py → branch_manager.py} +0 -0
- lionagi/core/{messages/messages.py → messages.py} +3 -3
- /lionagi/{_services → services}/__init__.py +0 -0
- /lionagi/{_services → services}/mistralai.py +0 -0
- /lionagi/{_services → services}/mlx_service.py +0 -0
- /lionagi/{_services → services}/ollama.py +0 -0
- /lionagi/{_services → services}/services.py +0 -0
- /lionagi/{_services → services}/transformers.py +0 -0
- {lionagi-0.0.208.dist-info → lionagi-0.0.210.dist-info}/LICENSE +0 -0
- {lionagi-0.0.208.dist-info → lionagi-0.0.210.dist-info}/WHEEL +0 -0
- {lionagi-0.0.208.dist-info → lionagi-0.0.210.dist-info}/top_level.txt +0 -0
lionagi/schema/data_logger.py
DELETED
@@ -1,126 +0,0 @@
|
|
1
|
-
from collections import deque
|
2
|
-
from typing import Dict, Any
|
3
|
-
from ..utils.sys_util import get_timestamp, create_path, as_dict
|
4
|
-
from ..utils.io_util import IOUtil
|
5
|
-
|
6
|
-
|
7
|
-
class DataLogger:
|
8
|
-
"""
|
9
|
-
A class for logging data entries and exporting them as CSV files.
|
10
|
-
|
11
|
-
This class provides functionality to log data entries in a deque and
|
12
|
-
supports exporting the logged data to a CSV file. The DataLogger can
|
13
|
-
be configured to use a specific directory for saving files.
|
14
|
-
|
15
|
-
Attributes:
|
16
|
-
dir (Optional[str]):
|
17
|
-
The default directory where CSV files will be saved.
|
18
|
-
log (deque):
|
19
|
-
A deque object that stores the logged data entries.
|
20
|
-
|
21
|
-
Methods:
|
22
|
-
__call__:
|
23
|
-
Adds an entry to the log.
|
24
|
-
to_csv:
|
25
|
-
Exports the logged data to a CSV file and clears the log.
|
26
|
-
set_dir:
|
27
|
-
Sets the default directory for saving CSV files.
|
28
|
-
"""
|
29
|
-
|
30
|
-
def __init__(self, dir= None, log: list = None) -> None:
|
31
|
-
"""
|
32
|
-
Initializes the DataLogger with an optional directory and initial log.
|
33
|
-
|
34
|
-
Parameters:
|
35
|
-
dir (Optional[str]): The directory where CSV files will be saved. Defaults to None.
|
36
|
-
|
37
|
-
log (Optional[List]): An initial list of log entries. Defaults to an empty list.
|
38
|
-
"""
|
39
|
-
self.dir = dir
|
40
|
-
self.log = deque(log) if log else deque()
|
41
|
-
|
42
|
-
def add_entry(self, entry: Dict[str, Any], level: str = "INFO") -> None:
|
43
|
-
"""
|
44
|
-
Adds a new entry to the log with a timestamp and a log level.
|
45
|
-
|
46
|
-
Args:
|
47
|
-
entry (Dict[str, Any]): The data entry to be added to the log.
|
48
|
-
level (str): The log level for the entry (e.g., "INFO", "ERROR"). Defaults to "INFO".
|
49
|
-
"""
|
50
|
-
self.log.append({
|
51
|
-
"timestamp": get_timestamp(), "level": level, **as_dict(entry)
|
52
|
-
})
|
53
|
-
|
54
|
-
def set_dir(self, dir: str) -> None:
|
55
|
-
"""
|
56
|
-
Sets the default directory for saving CSV files.
|
57
|
-
|
58
|
-
Parameters:
|
59
|
-
dir (str): The directory to be set as the default for saving files.
|
60
|
-
"""
|
61
|
-
self.dir = dir
|
62
|
-
|
63
|
-
def to_csv(
|
64
|
-
self, filename: str,
|
65
|
-
file_exist_ok: bool = False,
|
66
|
-
timestamp = True,
|
67
|
-
time_prefix: bool = False,
|
68
|
-
verbose: bool = True,
|
69
|
-
clear = True
|
70
|
-
) -> None:
|
71
|
-
"""
|
72
|
-
Exports the logged data to a CSV file, using the provided utilities for path creation and timestamping.
|
73
|
-
|
74
|
-
Args:
|
75
|
-
filename (str): The name of the CSV file.
|
76
|
-
file_exist_ok (bool): If True, creates the directory for the file if it does not exist. Defaults to False.
|
77
|
-
verbose (bool): If True, prints a message upon completion. Defaults to True.
|
78
|
-
time_prefix (bool): If True, adds the timestamp as a prefix to the filename. Defaults to False.
|
79
|
-
"""
|
80
|
-
if not filename.endswith('.csv'):
|
81
|
-
filename += '.csv'
|
82
|
-
|
83
|
-
filepath = create_path(
|
84
|
-
self.dir, filename, timestamp=timestamp,
|
85
|
-
dir_exist_ok=file_exist_ok, time_prefix=time_prefix
|
86
|
-
)
|
87
|
-
IOUtil.to_csv(list(self.log), filepath)
|
88
|
-
|
89
|
-
if verbose:
|
90
|
-
print(f"{len(self.log)} logs saved to {filepath}")
|
91
|
-
|
92
|
-
if clear:
|
93
|
-
self.log.clear()
|
94
|
-
|
95
|
-
def to_jsonl(
|
96
|
-
self, filename: str,
|
97
|
-
timestamp = False,
|
98
|
-
time_prefix=False,
|
99
|
-
file_exist_ok: bool = False,
|
100
|
-
verbose: bool = True,
|
101
|
-
clear = True
|
102
|
-
) -> None:
|
103
|
-
"""
|
104
|
-
Exports the logged data to a JSONL file and optionally clears the log.
|
105
|
-
|
106
|
-
Parameters:
|
107
|
-
filename (str): The name of the JSONL file.
|
108
|
-
file_exist_ok (bool): If True, creates the directory for the file if it does not exist. Defaults to False.
|
109
|
-
verbose (bool): If True, prints a message upon completion. Defaults to True.
|
110
|
-
"""
|
111
|
-
if not filename.endswith('.jsonl'):
|
112
|
-
filename += '.jsonl'
|
113
|
-
|
114
|
-
filepath = create_path(
|
115
|
-
self.dir, filename, timestamp=timestamp,
|
116
|
-
dir_exist_ok=file_exist_ok, time_prefix=time_prefix
|
117
|
-
)
|
118
|
-
|
119
|
-
for entry in self.log:
|
120
|
-
IOUtil.append_to_jsonl(entry, filepath)
|
121
|
-
|
122
|
-
if verbose:
|
123
|
-
print(f"{len(self.log)} logs saved to {filepath}")
|
124
|
-
|
125
|
-
if clear:
|
126
|
-
self.log.clear()
|
lionagi/schema/data_node.py
DELETED
@@ -1,88 +0,0 @@
|
|
1
|
-
from .base_node import BaseNode
|
2
|
-
from typing import Any
|
3
|
-
|
4
|
-
|
5
|
-
class DataNode(BaseNode):
|
6
|
-
|
7
|
-
def to_llama_index(self, **kwargs) -> Any:
|
8
|
-
"""
|
9
|
-
Converts node to llama index format.
|
10
|
-
|
11
|
-
Args:
|
12
|
-
**kwargs: Variable length argument list.
|
13
|
-
|
14
|
-
Returns:
|
15
|
-
The llama index representation of the node.
|
16
|
-
|
17
|
-
Examples:
|
18
|
-
node = DataNode()
|
19
|
-
llama_index = node.to_llama_index()
|
20
|
-
"""
|
21
|
-
from lionagi.bridge.llama_index import to_llama_index_textnode
|
22
|
-
return to_llama_index_textnode(self, **kwargs)
|
23
|
-
|
24
|
-
def to_langchain(self, **kwargs) -> Any:
|
25
|
-
"""
|
26
|
-
Converts node to langchain document format.
|
27
|
-
|
28
|
-
Args:
|
29
|
-
**kwargs: Variable length argument list.
|
30
|
-
|
31
|
-
Returns:
|
32
|
-
The langchain document representation of the node.
|
33
|
-
|
34
|
-
Examples:
|
35
|
-
node = DataNode()
|
36
|
-
langchain_doc = node.to_langchain()
|
37
|
-
"""
|
38
|
-
from lionagi.bridge.langchain import to_langchain_document
|
39
|
-
return to_langchain_document(self, **kwargs)
|
40
|
-
|
41
|
-
@classmethod
|
42
|
-
def from_llama_index(cls, llama_node: Any, **kwargs) -> "DataNode":
|
43
|
-
"""
|
44
|
-
Creates a DataNode instance from a llama index node.
|
45
|
-
|
46
|
-
Args:
|
47
|
-
llama_node: The llama index node object.
|
48
|
-
**kwargs: Variable length argument list.
|
49
|
-
|
50
|
-
Returns:
|
51
|
-
An instance of DataNode.
|
52
|
-
|
53
|
-
Examples:
|
54
|
-
llama_node = SomeLlamaIndexNode()
|
55
|
-
data_node = DataNode.from_llama_index(llama_node)
|
56
|
-
"""
|
57
|
-
llama_dict = llama_node.to_dict(**kwargs)
|
58
|
-
return cls.from_dict(llama_dict)
|
59
|
-
|
60
|
-
@classmethod
|
61
|
-
def from_langchain(cls, lc_doc: Any) -> "DataNode":
|
62
|
-
"""
|
63
|
-
Creates a DataNode instance from a langchain document.
|
64
|
-
|
65
|
-
Args:
|
66
|
-
lc_doc: The langchain document object.
|
67
|
-
|
68
|
-
Returns:
|
69
|
-
An instance of DataNode.
|
70
|
-
|
71
|
-
Examples:
|
72
|
-
lc_doc = SomeLangChainDocument()
|
73
|
-
data_node = DataNode.from_langchain(lc_doc)
|
74
|
-
"""
|
75
|
-
info_json = lc_doc.to_json()
|
76
|
-
info_node = {'lc_id': info_json['id']}
|
77
|
-
info_node = {**info_node, **info_json['kwargs']}
|
78
|
-
return cls(**info_node)
|
79
|
-
|
80
|
-
|
81
|
-
class File(DataNode):
|
82
|
-
|
83
|
-
...
|
84
|
-
|
85
|
-
|
86
|
-
class Chunk(DataNode):
|
87
|
-
|
88
|
-
...
|
lionagi/schema/status_tracker.py
DELETED
@@ -1,37 +0,0 @@
|
|
1
|
-
from dataclasses import dataclass
|
2
|
-
|
3
|
-
|
4
|
-
# credit to OpenAI for the following object
|
5
|
-
@dataclass
|
6
|
-
class StatusTracker:
|
7
|
-
"""
|
8
|
-
Class for keeping track of various task statuses.
|
9
|
-
|
10
|
-
This class serves as a simple way to monitor different types of task
|
11
|
-
outcomes and errors within a system. It uses dataclasses for easy
|
12
|
-
creation and management of state.
|
13
|
-
|
14
|
-
Attributes:
|
15
|
-
num_tasks_started:
|
16
|
-
The number of tasks that have been initiated.
|
17
|
-
num_tasks_in_progress:
|
18
|
-
The number of tasks currently being processed.
|
19
|
-
num_tasks_succeeded:
|
20
|
-
The number of tasks that have completed successfully.
|
21
|
-
num_tasks_failed:
|
22
|
-
The number of tasks that have failed.
|
23
|
-
num_rate_limit_errors:
|
24
|
-
The number of tasks that failed due to rate limiting.
|
25
|
-
num_api_errors:
|
26
|
-
The number of tasks that failed due to API errors.
|
27
|
-
num_other_errors:
|
28
|
-
The number of tasks that failed due to other errors.
|
29
|
-
"""
|
30
|
-
num_tasks_started: int = 0
|
31
|
-
num_tasks_in_progress: int = 0
|
32
|
-
num_tasks_succeeded: int = 0
|
33
|
-
num_tasks_failed: int = 0
|
34
|
-
num_rate_limit_errors: int = 0
|
35
|
-
num_api_errors: int = 0
|
36
|
-
num_other_errors: int = 0
|
37
|
-
|
@@ -1,323 +0,0 @@
|
|
1
|
-
import unittest
|
2
|
-
import os
|
3
|
-
import tempfile
|
4
|
-
import zipfile
|
5
|
-
import hashlib
|
6
|
-
from cryptography.fernet import InvalidToken
|
7
|
-
from lionagi.utils.encrypt_util import EncrytionUtil
|
8
|
-
|
9
|
-
|
10
|
-
class TestPasswordStrengthChecker(unittest.TestCase):
|
11
|
-
|
12
|
-
def test_short_passwords(self):
|
13
|
-
self.assertFalse(EncrytionUtil.password_strength_checker("Short1"))
|
14
|
-
self.assertFalse(EncrytionUtil.password_strength_checker("A1"))
|
15
|
-
|
16
|
-
def test_passwords_without_digits(self):
|
17
|
-
self.assertFalse(EncrytionUtil.password_strength_checker("NoDigitsHere"))
|
18
|
-
self.assertFalse(EncrytionUtil.password_strength_checker("Onlyletters!"))
|
19
|
-
|
20
|
-
def test_passwords_without_uppercase(self):
|
21
|
-
self.assertFalse(EncrytionUtil.password_strength_checker("alllowercase1"))
|
22
|
-
self.assertFalse(EncrytionUtil.password_strength_checker("nouppercase1!"))
|
23
|
-
|
24
|
-
def test_strong_passwords(self):
|
25
|
-
self.assertTrue(EncrytionUtil.password_strength_checker("ValidPass1"))
|
26
|
-
self.assertTrue(EncrytionUtil.password_strength_checker("AnotherGood1"))
|
27
|
-
|
28
|
-
|
29
|
-
class TestGenerateEncryptionKey(unittest.TestCase):
|
30
|
-
|
31
|
-
def setUp(self):
|
32
|
-
# Strong password and predefined salt for testing
|
33
|
-
self.strong_password = "StrongPass1"
|
34
|
-
self.salt = b'0123456789abcdef'
|
35
|
-
|
36
|
-
def test_with_strong_password_and_provided_salt(self):
|
37
|
-
key = EncrytionUtil.generate_encryption_key(password=self.strong_password, salt=self.salt)
|
38
|
-
self.assertIsInstance(key, str)
|
39
|
-
|
40
|
-
def test_with_strong_password_and_no_salt(self):
|
41
|
-
key = EncrytionUtil.generate_encryption_key(password=self.strong_password)
|
42
|
-
self.assertIsInstance(key, str)
|
43
|
-
|
44
|
-
def test_with_weak_password(self):
|
45
|
-
with self.assertRaises(ValueError):
|
46
|
-
EncrytionUtil.generate_encryption_key(password="weak")
|
47
|
-
|
48
|
-
def test_with_no_password(self):
|
49
|
-
key = EncrytionUtil.generate_encryption_key()
|
50
|
-
self.assertIsInstance(key, str)
|
51
|
-
self.assertEqual(len(key), 44) # Typical length of a Fernet key
|
52
|
-
|
53
|
-
|
54
|
-
class TestEncrypt(unittest.TestCase):
|
55
|
-
|
56
|
-
def setUp(self):
|
57
|
-
self.valid_key = EncrytionUtil.generate_encryption_key("StrongPass1")
|
58
|
-
self.invalid_key = "invalidkey"
|
59
|
-
self.test_data = "This is a test string."
|
60
|
-
|
61
|
-
def test_valid_encryption(self):
|
62
|
-
encrypted_data = EncrytionUtil.encrypt(self.test_data, self.valid_key)
|
63
|
-
self.assertIsInstance(encrypted_data, str)
|
64
|
-
self.assertNotEqual(encrypted_data, self.test_data)
|
65
|
-
|
66
|
-
def test_with_invalid_key(self):
|
67
|
-
with self.assertRaises(Exception):
|
68
|
-
EncrytionUtil.encrypt(self.test_data, self.invalid_key)
|
69
|
-
|
70
|
-
def test_with_non_string_data(self):
|
71
|
-
with self.assertRaises(AttributeError): # or whichever error is appropriate
|
72
|
-
EncrytionUtil.encrypt(12345, self.valid_key)
|
73
|
-
|
74
|
-
def test_with_empty_string(self):
|
75
|
-
encrypted_data = EncrytionUtil.encrypt("", self.valid_key)
|
76
|
-
self.assertIsInstance(encrypted_data, str)
|
77
|
-
self.assertNotEqual(encrypted_data, "")
|
78
|
-
|
79
|
-
|
80
|
-
class TestDecrypt(unittest.TestCase):
|
81
|
-
|
82
|
-
def setUp(self):
|
83
|
-
self.valid_key = EncrytionUtil.generate_encryption_key("StrongPass1")
|
84
|
-
self.invalid_key = "invalidkey"
|
85
|
-
self.test_data = "This is a test string."
|
86
|
-
self.encrypted_data = EncrytionUtil.encrypt(self.test_data, self.valid_key)
|
87
|
-
|
88
|
-
def test_valid_decryption(self):
|
89
|
-
decrypted_data = EncrytionUtil.decrypt(self.encrypted_data, self.valid_key)
|
90
|
-
self.assertIsInstance(decrypted_data, str)
|
91
|
-
self.assertEqual(decrypted_data, self.test_data)
|
92
|
-
|
93
|
-
def test_decryption_with_invalid_key(self):
|
94
|
-
with self.assertRaises(ValueError):
|
95
|
-
EncrytionUtil.decrypt(self.encrypted_data, self.invalid_key)
|
96
|
-
|
97
|
-
def test_with_non_string_data(self):
|
98
|
-
with self.assertRaises(AttributeError):
|
99
|
-
EncrytionUtil.decrypt(12345, self.valid_key)
|
100
|
-
|
101
|
-
def test_decryption_of_non_encrypted_string(self):
|
102
|
-
with self.assertRaises(InvalidToken):
|
103
|
-
EncrytionUtil.decrypt("plain text", self.valid_key)
|
104
|
-
|
105
|
-
|
106
|
-
class TestEncryptFile(unittest.TestCase):
|
107
|
-
|
108
|
-
def setUp(self):
|
109
|
-
self.valid_key = EncrytionUtil.generate_encryption_key("StrongPass1")
|
110
|
-
# Create a temporary file with some test data
|
111
|
-
self.temp_file = tempfile.NamedTemporaryFile(delete=False)
|
112
|
-
self.temp_file.write(b"This is a test file.")
|
113
|
-
self.temp_file.close()
|
114
|
-
|
115
|
-
def tearDown(self):
|
116
|
-
# Cleanup: Remove temporary files
|
117
|
-
os.remove(self.temp_file.name)
|
118
|
-
if os.path.exists(self.temp_file.name + '.enc'):
|
119
|
-
os.remove(self.temp_file.name + '.enc')
|
120
|
-
|
121
|
-
def test_encrypting_valid_file(self):
|
122
|
-
EncrytionUtil.encrypt_file(self.temp_file.name, self.valid_key)
|
123
|
-
self.assertTrue(os.path.exists(self.temp_file.name + '.enc'))
|
124
|
-
|
125
|
-
def test_with_non_existent_file_path(self):
|
126
|
-
with self.assertRaises(FileNotFoundError):
|
127
|
-
EncrytionUtil.encrypt_file("non_existent_file.txt", self.valid_key)
|
128
|
-
|
129
|
-
|
130
|
-
class TestDecryptFile(unittest.TestCase):
|
131
|
-
|
132
|
-
def setUp(self):
|
133
|
-
self.valid_key = EncrytionUtil.generate_encryption_key("StrongPass1")
|
134
|
-
# Create a temporary file and encrypt it
|
135
|
-
self.temp_file = tempfile.NamedTemporaryFile(delete=False)
|
136
|
-
self.temp_file.write(b"This is a test file.")
|
137
|
-
self.temp_file.close()
|
138
|
-
EncrytionUtil.encrypt_file(self.temp_file.name, self.valid_key)
|
139
|
-
|
140
|
-
def tearDown(self):
|
141
|
-
# Cleanup: Remove temporary files
|
142
|
-
os.remove(self.temp_file.name)
|
143
|
-
if os.path.exists(self.temp_file.name + '.enc'):
|
144
|
-
os.remove(self.temp_file.name + '.enc')
|
145
|
-
decrypted_file_path = self.temp_file.name.replace('.enc', '')
|
146
|
-
if os.path.exists(decrypted_file_path):
|
147
|
-
os.remove(decrypted_file_path)
|
148
|
-
|
149
|
-
def test_decrypting_valid_encrypted_file(self):
|
150
|
-
EncrytionUtil.decrypt_file(self.temp_file.name + '.enc', self.valid_key)
|
151
|
-
decrypted_file_path = self.temp_file.name.replace('.enc', '')
|
152
|
-
self.assertTrue(os.path.exists(decrypted_file_path))
|
153
|
-
|
154
|
-
def test_with_non_existent_encrypted_file_path(self):
|
155
|
-
with self.assertRaises(FileNotFoundError):
|
156
|
-
EncrytionUtil.decrypt_file("non_existent_file.txt.enc", self.valid_key)
|
157
|
-
|
158
|
-
def test_decrypting_with_invalid_key(self):
|
159
|
-
with self.assertRaises(ValueError):
|
160
|
-
EncrytionUtil.decrypt_file(self.temp_file.name + '.enc', 'invalidkey')
|
161
|
-
|
162
|
-
|
163
|
-
class TestIsEncrypted(unittest.TestCase):
|
164
|
-
|
165
|
-
def setUp(self):
|
166
|
-
self.valid_key = EncrytionUtil.generate_encryption_key("StrongPass1")
|
167
|
-
# Create a temporary file and encrypt it
|
168
|
-
self.temp_file = tempfile.NamedTemporaryFile(delete=False)
|
169
|
-
self.temp_file.write(b"This is a test file.")
|
170
|
-
self.temp_file.close()
|
171
|
-
EncrytionUtil.encrypt_file(self.temp_file.name, self.valid_key)
|
172
|
-
|
173
|
-
# Create a non-encrypted temporary file
|
174
|
-
self.non_encrypted_file = tempfile.NamedTemporaryFile(delete=False)
|
175
|
-
self.non_encrypted_file.write(b"This is a non-encrypted test file.")
|
176
|
-
self.non_encrypted_file.close()
|
177
|
-
|
178
|
-
def tearDown(self):
|
179
|
-
# Cleanup: Remove temporary files
|
180
|
-
os.remove(self.temp_file.name)
|
181
|
-
if os.path.exists(self.temp_file.name + '.enc'):
|
182
|
-
os.remove(self.temp_file.name + '.enc')
|
183
|
-
os.remove(self.non_encrypted_file.name)
|
184
|
-
|
185
|
-
def test_with_encrypted_file(self):
|
186
|
-
self.assertTrue(EncrytionUtil.is_encrypted(self.temp_file.name + '.enc', self.valid_key))
|
187
|
-
|
188
|
-
def test_with_non_encrypted_file(self):
|
189
|
-
self.assertFalse(EncrytionUtil.is_encrypted(self.non_encrypted_file.name, self.valid_key))
|
190
|
-
|
191
|
-
|
192
|
-
class TestDecompressFile(unittest.TestCase):
|
193
|
-
|
194
|
-
def setUp(self):
|
195
|
-
# Create a temporary directory
|
196
|
-
self.temp_dir = tempfile.mkdtemp()
|
197
|
-
|
198
|
-
# Create a temporary zip file with some content
|
199
|
-
self.temp_zip_file = os.path.join(self.temp_dir, 'test.zip')
|
200
|
-
with zipfile.ZipFile(self.temp_zip_file, 'w') as zipf:
|
201
|
-
zipf.writestr('test.txt', 'This is a test file.')
|
202
|
-
|
203
|
-
def tearDown(self):
|
204
|
-
# Cleanup: Remove temporary directory and its contents
|
205
|
-
for root, dirs, files in os.walk(self.temp_dir, topdown=False):
|
206
|
-
for name in files:
|
207
|
-
os.remove(os.path.join(root, name))
|
208
|
-
for name in dirs:
|
209
|
-
os.rmdir(os.path.join(root, name))
|
210
|
-
os.rmdir(self.temp_dir)
|
211
|
-
|
212
|
-
def test_decompressing_valid_zip_file(self):
|
213
|
-
EncrytionUtil.decompress_file(self.temp_zip_file, self.temp_dir)
|
214
|
-
self.assertTrue(os.path.exists(os.path.join(self.temp_dir, 'test.txt')))
|
215
|
-
|
216
|
-
def test_with_non_existent_zip_file_path(self):
|
217
|
-
with self.assertRaises(FileNotFoundError):
|
218
|
-
EncrytionUtil.decompress_file("non_existent_file.zip", self.temp_dir)
|
219
|
-
|
220
|
-
def test_with_invalid_zip_file(self):
|
221
|
-
invalid_zip_file = os.path.join(self.temp_dir, 'invalid.zip')
|
222
|
-
with open(invalid_zip_file, 'w') as f:
|
223
|
-
f.write("This is not a zip file.")
|
224
|
-
with self.assertRaises(zipfile.BadZipFile):
|
225
|
-
EncrytionUtil.decompress_file(invalid_zip_file, self.temp_dir)
|
226
|
-
|
227
|
-
|
228
|
-
class TestCompressFile(unittest.TestCase):
|
229
|
-
|
230
|
-
def setUp(self):
|
231
|
-
# Create a temporary file with some test data
|
232
|
-
self.temp_file = tempfile.NamedTemporaryFile(delete=False)
|
233
|
-
self.temp_file.write(b"This is a test file.")
|
234
|
-
self.temp_file.close()
|
235
|
-
|
236
|
-
def tearDown(self):
|
237
|
-
# Cleanup: Remove temporary files
|
238
|
-
os.remove(self.temp_file.name)
|
239
|
-
if os.path.exists(self.temp_file.name + '.zip'):
|
240
|
-
os.remove(self.temp_file.name + '.zip')
|
241
|
-
|
242
|
-
def test_compressing_valid_file(self):
|
243
|
-
EncrytionUtil.compress_file(self.temp_file.name)
|
244
|
-
self.assertTrue(os.path.exists(self.temp_file.name + '.zip'))
|
245
|
-
|
246
|
-
def test_with_non_existent_file_path(self):
|
247
|
-
with self.assertRaises(FileNotFoundError):
|
248
|
-
EncrytionUtil.compress_file("non_existent_file.txt")
|
249
|
-
|
250
|
-
|
251
|
-
class TestBinaryToHex(unittest.TestCase):
|
252
|
-
|
253
|
-
def test_with_valid_binary_data(self):
|
254
|
-
# Test a variety of binary data
|
255
|
-
self.assertEqual(EncrytionUtil.binary_to_hex(b'\x00\x0F'), '000f')
|
256
|
-
self.assertEqual(EncrytionUtil.binary_to_hex(b'hello'), '68656c6c6f')
|
257
|
-
self.assertEqual(EncrytionUtil.binary_to_hex(b'\xff\xfe\xfd\xfc'), 'fffefdfc')
|
258
|
-
|
259
|
-
def test_with_empty_bytes(self):
|
260
|
-
self.assertEqual(EncrytionUtil.binary_to_hex(b''), '')
|
261
|
-
|
262
|
-
|
263
|
-
class TestCreateHash(unittest.TestCase):
|
264
|
-
|
265
|
-
def test_hashing_with_default_algorithm(self):
|
266
|
-
data = "test"
|
267
|
-
expected_hash = hashlib.sha256(data.encode()).hexdigest()
|
268
|
-
self.assertEqual(EncrytionUtil.create_hash(data), expected_hash)
|
269
|
-
|
270
|
-
def test_hashing_with_different_algorithms(self):
|
271
|
-
data = "test"
|
272
|
-
algorithms = ['sha1', 'sha224', 'sha384', 'sha512']
|
273
|
-
for algo in algorithms:
|
274
|
-
with self.subTest(algorithm=algo):
|
275
|
-
expected_hash = hashlib.new(algo, data.encode()).hexdigest()
|
276
|
-
self.assertEqual(EncrytionUtil.create_hash(data, algo), expected_hash)
|
277
|
-
|
278
|
-
def test_with_unsupported_algorithm(self):
|
279
|
-
with self.assertRaises(ValueError):
|
280
|
-
EncrytionUtil.create_hash("test", "unsupported_algo")
|
281
|
-
|
282
|
-
|
283
|
-
class TestDecodeBase64(unittest.TestCase):
|
284
|
-
|
285
|
-
def test_with_valid_base64_encoded_strings(self):
|
286
|
-
# Test a variety of valid base64 encoded strings
|
287
|
-
test_cases = [
|
288
|
-
("SGVsbG8sIFdvcmxkIQ==", "Hello, World!"),
|
289
|
-
("VGhpcyBpcyBhIHRlc3Q=", "This is a test"),
|
290
|
-
("c29tZSBieXRlcw==", "some bytes")
|
291
|
-
]
|
292
|
-
for encoded, original in test_cases:
|
293
|
-
with self.subTest(encoded=encoded):
|
294
|
-
self.assertEqual(EncrytionUtil.decode_base64(encoded), original)
|
295
|
-
|
296
|
-
def test_with_invalid_base64_string(self):
|
297
|
-
invalid_data = "This is not base64!"
|
298
|
-
with self.assertRaises(Exception): # Replace Exception with the specific exception if known
|
299
|
-
EncrytionUtil.decode_base64(invalid_data)
|
300
|
-
|
301
|
-
def test_with_empty_string(self):
|
302
|
-
self.assertEqual(EncrytionUtil.decode_base64(""), "")
|
303
|
-
|
304
|
-
|
305
|
-
class TestEncodeBase64(unittest.TestCase):
|
306
|
-
|
307
|
-
def test_with_valid_strings(self):
|
308
|
-
# Test a variety of strings
|
309
|
-
test_cases = [
|
310
|
-
("Hello, World!", "SGVsbG8sIFdvcmxkIQ=="),
|
311
|
-
("This is a test", "VGhpcyBpcyBhIHRlc3Q="),
|
312
|
-
("some bytes", "c29tZSBieXRlcw==")
|
313
|
-
]
|
314
|
-
for original, encoded in test_cases:
|
315
|
-
with self.subTest(original=original):
|
316
|
-
self.assertEqual(EncrytionUtil.encode_base64(original), encoded)
|
317
|
-
|
318
|
-
def test_with_empty_string(self):
|
319
|
-
self.assertEqual(EncrytionUtil.encode_base64(""), "")
|
320
|
-
|
321
|
-
|
322
|
-
if __name__ == '__main__':
|
323
|
-
unittest.main()
|