openai/openai-python
Publicmirrored from https://github.com/openai/openai-pythonAvailable
tests/test_extract_files.py
73lines · modecode
| 1 | from __future__ import annotations |
| 2 | |
| 3 | from typing import Sequence |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from openai._types import FileTypes |
| 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>"]]) == [ |
| 41 | ("files[]", b"file one"), |
| 42 | ("files[]", b"file two"), |
| 43 | ] |
| 44 | assert query == {"title": "hello"} |
| 45 | |
| 46 | |
| 47 | @pytest.mark.parametrize( |
| 48 | "query,paths,expected", |
| 49 | [ |
| 50 | [ |
| 51 | {"foo": {"bar": "baz"}}, |
| 52 | [["foo", "<array>", "bar"]], |
| 53 | [], |
| 54 | ], |
| 55 | [ |
| 56 | {"foo": ["bar", "baz"]}, |
| 57 | [["foo", "bar"]], |
| 58 | [], |
| 59 | ], |
| 60 | [ |
| 61 | {"foo": {"bar": "baz"}}, |
| 62 | [["foo", "foo"]], |
| 63 | [], |
| 64 | ], |
| 65 | ], |
| 66 | ids=["dict expecting array", "array expecting dict", "unknown keys"], |
| 67 | ) |
| 68 | def test_ignores_incorrect_paths( |
| 69 | query: dict[str, object], |
| 70 | paths: Sequence[Sequence[str]], |
| 71 | expected: list[tuple[str, FileTypes]], |
| 72 | ) -> None: |
| 73 | assert extract_files(query, paths=paths) == expected |
| 74 | |