Skip to content

Commit 60017a9

Browse files
authored
Add test for ffmpeg in base images (#2122)
* Add test for ffmpeg in base images * Add predict.py to ffmpeg test * Fix pip_freeze integration test
1 parent b7aa7c3 commit 60017a9

File tree

3 files changed

+35
-2
lines changed

3 files changed

+35
-2
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
build:
2+
gpu: true
3+
python_version: "3.9"
4+
python_packages:
5+
- "torch==2.5.1"
6+
cuda: "12.4"
7+
run:
8+
- command: ffmpeg --help
9+
predict: "predict.py:Predictor"
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
from cog import BasePredictor
2+
3+
4+
class Predictor(BasePredictor):
5+
def predict(self, s: str) -> str:
6+
return "hello " + s

test-integration/test_integration/test_build.py

+20-2
Original file line numberDiff line numberDiff line change
@@ -381,11 +381,17 @@ def test_pip_freeze(docker_image):
381381
labels = image[0]["Config"]["Labels"]
382382
pip_freeze = labels["run.cog.pip_freeze"]
383383
pip_freeze = "\n".join(
384-
[x for x in pip_freeze.split("\n") if not x.startswith("cog @")]
384+
[
385+
x
386+
for x in pip_freeze.split("\n")
387+
if not x.startswith("cog @")
388+
and not x.startswith("fastapi")
389+
and not x.startswith("starlette")
390+
]
385391
)
386392
assert (
387393
pip_freeze
388-
== "anyio==4.4.0\nattrs==23.2.0\ncertifi==2024.8.30\ncharset-normalizer==3.3.2\nclick==8.1.7\nexceptiongroup==1.2.2\nfastapi==0.115.6\nh11==0.14.0\nhttptools==0.6.1\nidna==3.8\npydantic==1.10.18\npython-dotenv==1.0.1\nPyYAML==6.0.2\nrequests==2.32.3\nsniffio==1.3.1\nstarlette==0.41.3\nstructlog==24.4.0\ntyping_extensions==4.12.2\nurllib3==2.2.2\nuvicorn==0.30.6\nuvloop==0.20.0\nwatchfiles==0.24.0\nwebsockets==13.0.1\n"
394+
== "anyio==4.4.0\nattrs==23.2.0\ncertifi==2024.8.30\ncharset-normalizer==3.3.2\nclick==8.1.7\nexceptiongroup==1.2.2\nh11==0.14.0\nhttptools==0.6.1\nidna==3.8\npydantic==1.10.18\npython-dotenv==1.0.1\nPyYAML==6.0.2\nrequests==2.32.3\nsniffio==1.3.1\nstructlog==24.4.0\ntyping_extensions==4.12.2\nurllib3==2.2.2\nuvicorn==0.30.6\nuvloop==0.20.0\nwatchfiles==0.24.0\nwebsockets==13.0.1\n"
389395
)
390396

391397

@@ -434,3 +440,15 @@ def test_pydantic2(docker_image):
434440
)
435441

436442
assert build_process.returncode == 0
443+
444+
445+
def test_ffmpeg_base_image(docker_image):
446+
project_dir = Path(__file__).parent / "fixtures/ffmpeg-package"
447+
448+
build_process = subprocess.run(
449+
["cog", "build", "-t", docker_image],
450+
cwd=project_dir,
451+
capture_output=True,
452+
)
453+
454+
assert build_process.returncode == 0

0 commit comments

Comments
 (0)