openai/openai-python
Publicmirrored fromhttps://github.com/openai/openai-pythonAvailable
tests/test_extract_files.py
91lines · modecode
| 1 | from __future__ import annotations |
| 2 | |
| 3 | from typing import Sequence |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from openai._types import FileTypes, ArrayFormat |
| 8 | from openai._utils import extract_files |
| 9 | |
| 10 | |
| 11 | def test_removes_files_from_input() -> None: |
| 12 | query = {"foo": "bar"} |
| 13 | assert extract_files(query, paths=[]) == [] |
| 14 | assert query == {"foo": "bar"} |
| 15 | |
| 16 | query2 = {"foo": b"Bar", "hello": "world"} |
| 17 | assert extract_files(query2, paths=[["foo"]]) == [("foo", b"Bar")] |
| 18 | assert query2 == {"hello": "world"} |
| 19 | |
| 20 | query3 = {"foo": {"foo": {"bar": b"Bar"}}, "hello": "world"} |
| 21 | assert extract_files(query3, paths=[["foo", "foo", "bar"]]) == [("foo[foo][bar]", b"Bar")] |
| 22 | assert query3 == {"foo": {"foo": {}}, "hello": "world"} |
| 23 | |
| 24 | query4 = {"foo": {"bar": b"Bar", "baz": "foo"}, "hello": "world"} |
| 25 | assert extract_files(query4, paths=[["foo", "bar"]]) == [("foo[bar]", b"Bar")] |
| 26 | assert query4 == {"hello": "world", "foo": {"baz": "foo"}} |
| 27 | |
| 28 | |
| 29 | def test_multiple_files() -> None: |
| 30 | query = {"documents": [{"file": b"My first file"}, {"file": b"My second file"}]} |
| 31 | assert extract_files(query, paths=[["documents", "<array>", "file"]]) == [ |
| 32 | ("documents[][file]", b"My first file"), |
| 33 | ("documents[][file]", b"My second file"), |
| 34 | ] |
| 35 | assert query == {"documents": [{}, {}]} |
| 36 | |
| 37 | |
| 38 | def test_top_level_file_array() -> None: |
| 39 | query = {"files": [b"file one", b"file two"], "title": "hello"} |
| 40 | assert extract_files(query, paths=[["files", "<array>"]]) == [("files[]", b"file one"), ("files[]", b"file two")] |
| 41 | assert query == {"title": "hello"} |
| 42 | |
| 43 | |
| 44 | @pytest.mark.parametrize( |
| 45 | "query,paths,expected", |
| 46 | [ |
| 47 | [ |
| 48 | {"foo": {"bar": "baz"}}, |
| 49 | [["foo", "<array>", "bar"]], |
| 50 | [], |
| 51 | ], |
| 52 | [ |
| 53 | {"foo": ["bar", "baz"]}, |
| 54 | [["foo", "bar"]], |
| 55 | [], |
| 56 | ], |
| 57 | [ |
| 58 | {"foo": {"bar": "baz"}}, |
| 59 | [["foo", "foo"]], |
| 60 | [], |
| 61 | ], |
| 62 | ], |
| 63 | ids=["dict expecting array", "array expecting dict", "unknown keys"], |
| 64 | ) |
| 65 | def test_ignores_incorrect_paths( |
| 66 | query: dict[str, object], |
| 67 | paths: Sequence[Sequence[str]], |
| 68 | expected: list[tuple[str, FileTypes]], |
| 69 | ) -> None: |
| 70 | assert extract_files(query, paths=paths) == expected |
| 71 | |
| 72 | |
| 73 | @pytest.mark.parametrize( |
| 74 | "array_format,expected_top_level,expected_nested", |
| 75 | [ |
| 76 | ("brackets", [("files[]", b"a"), ("files[]", b"b")], [("items[][file]", b"a"), ("items[][file]", b"b")]), |
| 77 | ("repeat", [("files", b"a"), ("files", b"b")], [("items[file]", b"a"), ("items[file]", b"b")]), |
| 78 | ("comma", [("files", b"a"), ("files", b"b")], [("items[file]", b"a"), ("items[file]", b"b")]), |
| 79 | ("indices", [("files[0]", b"a"), ("files[1]", b"b")], [("items[0][file]", b"a"), ("items[1][file]", b"b")]), |
| 80 | ], |
| 81 | ) |
| 82 | def test_array_format_controls_file_field_names( |
| 83 | array_format: ArrayFormat, |
| 84 | expected_top_level: list[tuple[str, FileTypes]], |
| 85 | expected_nested: list[tuple[str, FileTypes]], |
| 86 | ) -> None: |
| 87 | top_level = {"files": [b"a", b"b"]} |
| 88 | assert extract_files(top_level, paths=[["files", "<array>"]], array_format=array_format) == expected_top_level |
| 89 | |
| 90 | nested = {"items": [{"file": b"a"}, {"file": b"b"}]} |
| 91 | assert extract_files(nested, paths=[["items", "<array>", "file"]], array_format=array_format) == expected_nested |
| 92 | |