Repository: langgenius/dify
File: api/tests/unit_tests/controllers/console/test_document_detail_api_data_source_info.py (Python)

Code before:
    # Allow optional metadata occasionally present in tests
    # (kept optional so static checkers don't fail when it appears)
    created_at: NotRequired[str]


class UploadFileInfo(TypedDict):
    upload_file_id: str


class NotionImportInfo(TypedDict):
    notion_page_id: str
    workspace_id: str


class WebsiteCrawlInfo(TypedDict):
    url: str
    job_id: str


RawInfo = Union[LocalFileInfo, UploadFileInfo, NotionImportInfo, WebsiteCrawlInfo]


class Case(TypedDict):
    data_source_type: Literal["local_file", "upload_file", "notion_import", "website_crawl"]
    data_source_info: str
    expected_raw: RawInfo


class TestDocumentDetailDataSourceInfo:
    """Test cases for document detail API data_source_info serialization."""

    def test_data_source_info_dict_returns_raw_data(self):
        """Test that data_source_info_dict returns raw JSON data for all data_source_type values."""
        # Test data for different data_source_type values
        test_cases: list[Case] = [
            {
                "data_source_type": "local_file",
                "data_source_info": json.dumps({"file_path": "/tmp/test.txt", "size": 1024}),
                "expected_raw": {"file_path": "/tmp/test.txt", "size": 1024},
            },
            {
                "data_source_type": "upload_file",
                "data_source_info": json.dumps({"upload_file_id": "test-file-id"}),
                "expected_raw": {"upload_file_id": "test-file-id"},
            },
            {
                "data_source_type": "notion_import",
                "data_source_info": json.dumps({"notion_page_id": "page-123", "workspace_id": "ws-456"}),
                "expected_raw": {"notion_page_id": "page-123", "workspace_id": "ws-456"},
            },
            {

Review comment (suggestion):
```py
from typing import TypedDict, Literal, Union, Generic, TypeVar, NotRequired

# 1. Keep your existing Info definitions
class LocalFileInfo(TypedDict):
    file_path: str
    size: int
    created_at: NotRequired[str]

class UploadFileInfo(TypedDict):
    upload_file_id: str

class NotionImportInfo(TypedDict):
    notion_page_id: str
    workspace_id: str

class WebsiteCrawlInfo(TypedDict):
    url: str
    job_id: str

# 2. Define TypeVars for the Tag (Literal) and the Data (Info)
T_type = TypeVar("T_type", bound=str)
T_info = TypeVar("T_info", bound=Union[LocalFileInfo, UploadFileInfo, NotionImportInfo, WebsiteCrawlInfo])

# 3. Create the Generic TypedDict
class Case(TypedDict, Generic[T_type, T_info]):
    data_source_type: T_type
    data_source_info: str
    expected_raw: T_info

# 4. (Optional but recommended) Define a Union of the allowed specialized types
# This allows you to create a list of cases that are strictly validated
LocalFileCase = Case[Literal["local_file"], LocalFileInfo]
UploadFileCase = Case[Literal["upload_file"], UploadFileInfo]
NotionImportCase = Case[Literal["notion_import"], NotionImportInfo]
WebsiteCrawlCase = Case[Literal["website_crawl"], WebsiteCrawlInfo]

AnyCase = Union[LocalFileCase, UploadFileCase, NotionImportCase, WebsiteCrawlCase]

# This works perfectly
case_1: LocalFileCase = {
    "data_source_type": "local_file",
    "data_source_info": "some_string",
    "expected_raw": {
        "file_path": "/tmp/test.txt",
        "size": 1024
    }
}


# ERROR: Expected LocalFileInfo, but got WebsiteCrawlInfo
case_2: LocalFileCase = {
    "data_source_type": "local_file",
    "data_source_info": "...",
    "expected_raw": {
        "url": "https://google.com", 
        "job_id": "123"
    }
}

cases: list[AnyCase] = [case_1]

for case in cases:
    if case["data_source_type"] == "local_file":
        # Inside this block, the IDE knows 'expected_raw' has 'file_path'
        print(case["expected_raw"]["file_path"])
```

since 3.11 , we can use Generic to do type check.


Code after:

class NotionImportInfo(TypedDict):
    notion_page_id: str
    workspace_id: str


class WebsiteCrawlInfo(TypedDict):
    url: str
    job_id: str


RawInfo = Union[LocalFileInfo, UploadFileInfo, NotionImportInfo, WebsiteCrawlInfo]
T_type = TypeVar("T_type", bound=str)
T_info = TypeVar("T_info", bound=Union[LocalFileInfo, UploadFileInfo, NotionImportInfo, WebsiteCrawlInfo])


class Case(TypedDict, Generic[T_type, T_info]):
    data_source_type: T_type
    data_source_info: str
    expected_raw: T_info


LocalFileCase = Case[Literal["local_file"], LocalFileInfo]
UploadFileCase = Case[Literal["upload_file"], UploadFileInfo]
NotionImportCase = Case[Literal["notion_import"], NotionImportInfo]
WebsiteCrawlCase = Case[Literal["website_crawl"], WebsiteCrawlInfo]

AnyCase = Union[LocalFileCase, UploadFileCase, NotionImportCase, WebsiteCrawlCase]


case_1: LocalFileCase = {
    "data_source_type": "local_file",
    "data_source_info": json.dumps({"file_path": "/tmp/test.txt", "size": 1024}),
    "expected_raw": {"file_path": "/tmp/test.txt", "size": 1024},
}


# ERROR: Expected LocalFileInfo, but got WebsiteCrawlInfo
case_2: LocalFileCase = {
    "data_source_type": "local_file",
    "data_source_info": "...",
    "expected_raw": {"file_path": "https://google.com", "size": 123},
}

cases: list[AnyCase] = [case_1]


class TestDocumentDetailDataSourceInfo:
    """Test cases for document detail API data_source_info serialization."""

    def test_data_source_info_dict_returns_raw_data(self):