|
1 | 1 | import json |
| 2 | +import os |
2 | 3 | import threading |
3 | 4 | import time |
4 | 5 |
|
@@ -418,3 +419,78 @@ def test_missing_webbrowser_methods(fig1): |
418 | 419 | finally: |
419 | 420 | # restore everything after this test |
420 | 421 | webbrowser.get = removed_webbrowser_get_method |
| 422 | + |
| 423 | + |
| 424 | +def test_colab_renderer_when_env_var_is_set(): |
| 425 | + """ |
| 426 | + When COLAB_NOTEBOOK_ID is present the default renderer should be 'colab'. |
| 427 | + """ |
| 428 | + import importlib |
| 429 | + import plotly.io._renderers as _renderers_mod |
| 430 | + from plotly import optional_imports |
| 431 | + |
| 432 | + fake_ipython = MagicMock() |
| 433 | + original_get_module = optional_imports.get_module |
| 434 | + |
| 435 | + def patched_get_module(name, *args, **kwargs): |
| 436 | + if name in ("IPython", "IPython.display"): |
| 437 | + return fake_ipython |
| 438 | + return original_get_module(name, *args, **kwargs) |
| 439 | + |
| 440 | + original_default = pio.renderers.default |
| 441 | + try: |
| 442 | + with mock.patch.dict(os.environ, {"COLAB_NOTEBOOK_ID": "fake-id"}, clear=True): |
| 443 | + with mock.patch.object( |
| 444 | + optional_imports, "get_module", side_effect=patched_get_module |
| 445 | + ): |
| 446 | + importlib.reload(_renderers_mod) |
| 447 | + assert _renderers_mod.renderers.default == "colab" |
| 448 | + finally: |
| 449 | + importlib.reload(_renderers_mod) |
| 450 | + pio.renderers.default = original_default |
| 451 | + |
| 452 | + |
| 453 | +def test_colab_renderer_when_env_var_is_absent(): |
| 454 | + """ |
| 455 | + Without COLAB_NOTEBOOK_ID the default renderer must not be 'colab', |
| 456 | + even when ``google.colab`` is importable. The old detection used |
| 457 | + ``import google.colab`` which wrongly activated the colab renderer |
| 458 | + in environments that merely had the package installed (e.g. the |
| 459 | + Colab VS Code extension). |
| 460 | +
|
| 461 | + Regression test for https://github.com/plotly/plotly.py/pull/5473. |
| 462 | + """ |
| 463 | + import sys |
| 464 | + import types |
| 465 | + import importlib |
| 466 | + import plotly.io._renderers as _renderers_mod |
| 467 | + from plotly import optional_imports |
| 468 | + |
| 469 | + fake_ipython = MagicMock() |
| 470 | + original_get_module = optional_imports.get_module |
| 471 | + |
| 472 | + def patched_get_module(name, *args, **kwargs): |
| 473 | + if name in ("IPython", "IPython.display"): |
| 474 | + return fake_ipython |
| 475 | + return original_get_module(name, *args, **kwargs) |
| 476 | + |
| 477 | + # Make google.colab importable so the old ``import google.colab`` |
| 478 | + # approach would have chosen the colab renderer here. |
| 479 | + fake_google = types.ModuleType("google") |
| 480 | + fake_google_colab = types.ModuleType("google.colab") |
| 481 | + |
| 482 | + original_default = pio.renderers.default |
| 483 | + try: |
| 484 | + with mock.patch.dict(os.environ, {}, clear=True): |
| 485 | + with mock.patch.dict( |
| 486 | + sys.modules, |
| 487 | + {"google": fake_google, "google.colab": fake_google_colab}, |
| 488 | + ): |
| 489 | + with mock.patch.object( |
| 490 | + optional_imports, "get_module", side_effect=patched_get_module |
| 491 | + ): |
| 492 | + importlib.reload(_renderers_mod) |
| 493 | + assert _renderers_mod.renderers.default != "colab" |
| 494 | + finally: |
| 495 | + importlib.reload(_renderers_mod) |
| 496 | + pio.renderers.default = original_default |
0 commit comments