|
| 1 | +import json |
| 2 | +import logging |
| 3 | +import os |
| 4 | +import unittest |
| 5 | +from typing import Any, Optional |
| 6 | +from unittest.mock import Mock, patch |
| 7 | + |
| 8 | +from dotenv import load_dotenv |
| 9 | +from parameterized import parameterized |
| 10 | + |
| 11 | +from unstract.sdk.index import Index |
| 12 | +from unstract.sdk.tool.base import BaseTool |
| 13 | + |
| 14 | +load_dotenv() |
| 15 | + |
| 16 | +logger = logging.getLogger(__name__) |
| 17 | + |
| 18 | + |
| 19 | +def get_test_values(env_key: str) -> list[str]: |
| 20 | + test_values = json.loads(os.environ.get(env_key)) |
| 21 | + return test_values |
| 22 | + |
| 23 | + |
| 24 | +class ToolLLMTest(unittest.TestCase): |
| 25 | + class MockTool(BaseTool): |
| 26 | + def run( |
| 27 | + self, |
| 28 | + params: dict[str, Any] = {}, |
| 29 | + settings: dict[str, Any] = {}, |
| 30 | + workflow_id: str = "", |
| 31 | + ) -> None: |
| 32 | + # self.stream_log("Mock tool running") |
| 33 | + pass |
| 34 | + |
| 35 | + @classmethod |
| 36 | + def setUpClass(cls): |
| 37 | + cls.tool = cls.MockTool() |
| 38 | + |
| 39 | + @patch( |
| 40 | + "unstract.sdk.index.Index.generate_index_key", |
| 41 | + Mock( |
| 42 | + return_value="77843eb8d9e30ad56bfcb018c2633fa32feef2f0c09762b6b820c75664b64c1b" |
| 43 | + ), |
| 44 | + ) |
| 45 | + def test_generate_file_id(self): |
| 46 | + expected = "77843eb8d9e30ad56bfcb018c2633fa32feef2f0c09762b6b820c75664b64c1b" |
| 47 | + index = Index(tool=self.tool) |
| 48 | + actual = index.generate_file_id( |
| 49 | + tool_id="8ac26867-7811-4dc7-a17b-b16d3b561583", |
| 50 | + vector_db="81f1f6a8-cae8-4b8e-b2a4-57f80de512f6", |
| 51 | + embedding="ecf998d6-ded0-4aca-acd1-372a21daf0f9", |
| 52 | + x2text="59bc55fc-e2a7-48dd-ae93-794b4d81d46e", |
| 53 | + chunk_size="1024", |
| 54 | + chunk_overlap="128", |
| 55 | + file_path="/a/b/c", |
| 56 | + file_hash="045b1a67824592b67426f8e60c1f8328e8d2a35139f9983e0aa0a7b6f10915c3", |
| 57 | + ) |
| 58 | + assert expected == actual |
| 59 | + |
| 60 | + test_data = [ |
| 61 | + { |
| 62 | + "vector_db": "81f1f6a8-cae8-4b8e-b2a4-57f80de512f6", |
| 63 | + "embedding": "ecf998d6-ded0-4aca-acd1-372a21daf0f9", |
| 64 | + "x2text": "59bc55fc-e2a7-48dd-ae93-794b4d81d46e", |
| 65 | + "chunk_size": "1024", |
| 66 | + "chunk_overlap": "128", |
| 67 | + "file_path": "/a/b/c", |
| 68 | + }, |
| 69 | + { |
| 70 | + "vector_db": "81f1f6a8-cae8-4b8e-b2a4-57f80de512f6", |
| 71 | + "embedding": "ecf998d6-ded0-4aca-acd1-372a21daf0f9", |
| 72 | + "x2text": "59bc55fc-e2a7-48dd-ae93-794b4d81d46e", |
| 73 | + "chunk_size": "1024", |
| 74 | + "chunk_overlap": "128", |
| 75 | + "file_path": "/a/b/c", |
| 76 | + "file_hash": "045b1a67824592b67426f8e60c1f8328e8d2a35139f9983e0aa0a7b6f10915c3", |
| 77 | + }, |
| 78 | + ] |
| 79 | + |
| 80 | + @parameterized.expand(test_data) |
| 81 | + @patch( |
| 82 | + "unstract.sdk.adapter.ToolAdapter.get_adapter_config", |
| 83 | + Mock(return_value={}), |
| 84 | + ) |
| 85 | + @patch( |
| 86 | + "unstract.sdk.utils.ToolUtils.hash_str", |
| 87 | + Mock( |
| 88 | + return_value="77843eb8d9e30ad56bfcb018c2633fa32feef2f0c09762b6b820c75664b64c1b" |
| 89 | + ), |
| 90 | + ) |
| 91 | + @patch( |
| 92 | + "unstract.sdk.utils.ToolUtils.get_hash_from_file", |
| 93 | + Mock( |
| 94 | + return_value="ab940bb34a60d2a7876dd8e1bd1b22f5dc85936a9e2af3c49bfc888a1d944ff0" |
| 95 | + ), |
| 96 | + ) |
| 97 | + def test_generate_index_key( |
| 98 | + self, |
| 99 | + vector_db: str, |
| 100 | + embedding: str, |
| 101 | + x2text: str, |
| 102 | + chunk_size: str, |
| 103 | + chunk_overlap: str, |
| 104 | + file_path: Optional[str] = None, |
| 105 | + file_hash: Optional[str] = None, |
| 106 | + ): |
| 107 | + expected = "77843eb8d9e30ad56bfcb018c2633fa32feef2f0c09762b6b820c75664b64c1b" |
| 108 | + index = Index(tool=self.tool) |
| 109 | + actual = index.generate_index_key( |
| 110 | + vector_db, |
| 111 | + embedding, |
| 112 | + x2text, |
| 113 | + chunk_size, |
| 114 | + chunk_overlap, |
| 115 | + file_path, |
| 116 | + file_hash, |
| 117 | + ) |
| 118 | + assert expected == actual |
| 119 | + |
| 120 | + |
| 121 | +if __name__ == "__main__": |
| 122 | + unittest.main() |
0 commit comments