Fix memoir image prompt parsing and host allowlist
This commit is contained in:
@@ -176,6 +176,23 @@ class PollUntilCompleteTest(unittest.TestCase):
|
||||
|
||||
|
||||
class DownloadImageTest(unittest.TestCase):
|
||||
def test_download_allows_liblib_tmp_image_host_by_default(self):
|
||||
http_client = Mock()
|
||||
resp = Mock()
|
||||
resp.content = b"png-bytes"
|
||||
resp.raise_for_status = Mock()
|
||||
http_client.get.return_value = resp
|
||||
|
||||
provider = _make_provider(http_client)
|
||||
payload = provider.download_image(
|
||||
{"image_url": "https://liblibai-tmp-image.liblib.cloud/img/demo.png"}
|
||||
)
|
||||
|
||||
self.assertEqual(payload, b"png-bytes")
|
||||
http_client.get.assert_called_once_with(
|
||||
"https://liblibai-tmp-image.liblib.cloud/img/demo.png"
|
||||
)
|
||||
|
||||
@patch.dict(os.environ, {"MEMOIR_IMAGE_DOWNLOAD_HOSTS": "cdn.example.com"}, clear=False)
|
||||
def test_download_fetches_binary_payload(self):
|
||||
http_client = Mock()
|
||||
|
||||
Reference in New Issue
Block a user