code-puppy-core-plugins 0.0.2__py3-none-any.whl

code_puppy_core_plugins/chatgpt_oauth/test_plugin.py

@@ -0,0 +1,663 @@
+"""
+Basic tests for ChatGPT OAuth plugin.
+"""
+
+import asyncio
+import json
+from contextlib import nullcontext
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from code_puppy_core_plugins.chatgpt_oauth import (
+    config,
+    image_generation,
+    image_tool,
+    utils,
+)
+
+
+def test_codex_image_generation_posts_codex_payload(tmp_path):
+    encoded = "aGVsbG8="
+    response = MagicMock()
+    response.json.return_value = {"data": [{"b64_json": encoded}]}
+    response.raise_for_status.return_value = None
+
+    with (
+        patch.object(image_generation, "get_valid_access_token", return_value="token"),
+        patch.object(
+            image_generation,
+            "load_stored_tokens",
+            return_value={"account_id": "account"},
+        ),
+        patch.object(image_generation.config, "DATA_DIR", str(tmp_path)),
+        patch.object(image_generation.requests, "post", return_value=response) as post,
+    ):
+        output_path = image_generation.generate_image("a red fox")
+
+    assert output_path.read_bytes() == b"hello"
+    assert output_path.parent == tmp_path / "generated_images"
+    post.assert_called_once()
+    call = post.call_args
+    assert call.args[0].endswith("/images/generations")
+    assert call.kwargs["json"] == {
+        "prompt": "a red fox",
+        "background": "auto",
+        "model": "gpt-image-2",
+        "quality": "auto",
+        "size": "auto",
+    }
+    assert call.kwargs["headers"]["ChatGPT-Account-Id"] == "account"
+
+
+def test_codex_image_generation_with_references_posts_edits_payload(tmp_path):
+    """Reference images must go to /images/edits as [{"image_url": data-url}].
+
+    The Codex backend rejects every other shape with HTTP 400 (see the
+    image_generation module docstring), so this asserts the exact wire format.
+    """
+    reference = tmp_path / "ref.png"
+    reference.write_bytes(b"pngbytes")
+    expected_b64 = image_generation.base64.b64encode(b"pngbytes").decode()
+
+    encoded = "aGVsbG8="
+    response = MagicMock()
+    response.json.return_value = {"data": [{"b64_json": encoded}]}
+    response.raise_for_status.return_value = None
+
+    with (
+        patch.object(image_generation, "get_valid_access_token", return_value="token"),
+        patch.object(
+            image_generation,
+            "load_stored_tokens",
+            return_value={"account_id": "account"},
+        ),
+        patch.object(image_generation.config, "DATA_DIR", str(tmp_path)),
+        patch.object(image_generation.requests, "post", return_value=response) as post,
+    ):
+        image_generation.generate_image("a red fox", [reference])
+
+    call = post.call_args
+    assert call.args[0].endswith("/images/edits")
+    payload = call.kwargs["json"]
+    assert payload["images"] == [{"image_url": f"data:image/png;base64,{expected_b64}"}]
+    assert payload["prompt"] == "a red fox"
+    assert payload["model"] == "gpt-image-2"
+    # `background` is not accepted by the edits endpoint.
+    assert "background" not in payload
+
+
+def test_codex_image_generation_without_references_still_uses_generations(tmp_path):
+    """An empty/omitted reference list must not change existing behavior."""
+    response = MagicMock()
+    response.json.return_value = {"data": [{"b64_json": "aGVsbG8="}]}
+    response.raise_for_status.return_value = None
+
+    with (
+        patch.object(image_generation, "get_valid_access_token", return_value="token"),
+        patch.object(
+            image_generation,
+            "load_stored_tokens",
+            return_value={"account_id": "account"},
+        ),
+        patch.object(image_generation.config, "DATA_DIR", str(tmp_path)),
+        patch.object(image_generation.requests, "post", return_value=response) as post,
+    ):
+        image_generation.generate_image("a red fox", [])
+
+    assert post.call_args.args[0].endswith("/images/generations")
+    assert "images" not in post.call_args.kwargs["json"]
+
+
+def test_codex_image_generation_rejects_missing_reference(tmp_path):
+    with (
+        patch.object(image_generation, "get_valid_access_token", return_value="token"),
+        patch.object(
+            image_generation,
+            "load_stored_tokens",
+            return_value={"account_id": "account"},
+        ),
+        patch.object(image_generation.requests, "post") as post,
+        pytest.raises(image_generation.CodexImageGenerationError),
+    ):
+        image_generation.generate_image("a red fox", [tmp_path / "nope.png"])
+
+    post.assert_not_called()
+
+
+def test_codex_image_generation_rejects_oversized_reference(tmp_path):
+    reference = tmp_path / "huge.png"
+    reference.write_bytes(b"x" * 16)
+
+    with (
+        patch.object(image_generation, "get_valid_access_token", return_value="token"),
+        patch.object(
+            image_generation,
+            "load_stored_tokens",
+            return_value={"account_id": "account"},
+        ),
+        patch.object(image_generation, "_MAX_REFERENCE_BYTES", 8),
+        patch.object(image_generation.requests, "post") as post,
+        pytest.raises(image_generation.CodexImageGenerationError),
+    ):
+        image_generation.generate_image("a red fox", [reference])
+
+    post.assert_not_called()
+
+
+def test_codex_image_generation_requires_auth():
+    with (
+        patch.object(image_generation, "get_valid_access_token", return_value=None),
+        patch.object(image_generation, "load_stored_tokens", return_value={}),
+        pytest.raises(image_generation.CodexImageGenerationError, match="codex-auth"),
+    ):
+        image_generation.generate_image("a red fox")
+
+
+def test_emit_iterm_image(tmp_path):
+    output_path = tmp_path / "puppy.png"
+    output_path.write_bytes(b"image-bytes")
+    stdout = MagicMock()
+    stdout.isatty.return_value = True
+
+    with (
+        patch.dict(
+            image_generation.os.environ,
+            {"TERM_PROGRAM": "iTerm.app"},
+            clear=True,
+        ),
+        patch.object(image_generation.sys, "stdout", stdout),
+        patch(
+            "code_puppy.messaging.run_ui.suspended_run_ui",
+            return_value=nullcontext(),
+        ),
+    ):
+        assert image_generation.emit_iterm_image(output_path) is True
+
+    sequence = stdout.write.call_args.args[0]
+    assert sequence.startswith("\033]1337;File=name=")
+    assert ";inline=1;preserveAspectRatio=1:" in sequence
+    assert image_generation.base64.b64encode(b"image-bytes").decode() in sequence
+    stdout.flush.assert_called_once()
+
+
+def test_emit_iterm_image_ignores_other_terminals(tmp_path):
+    stdout = MagicMock()
+    stdout.isatty.return_value = True
+    with (
+        patch.dict(image_generation.os.environ, {}, clear=True),
+        patch.object(image_generation.sys, "stdout", stdout),
+    ):
+        assert image_generation.emit_iterm_image(tmp_path / "missing.png") is False
+    stdout.write.assert_not_called()
+
+
+def test_codex_imagegen_command():
+    from code_puppy_core_plugins.chatgpt_oauth import register_callbacks
+
+    with (
+        patch.object(
+            image_generation,
+            "generate_image",
+            return_value=Path("/tmp/generated.png"),
+        ) as generate,
+        patch.object(register_callbacks, "emit_info"),
+        patch.object(register_callbacks, "emit_success") as success,
+        patch.object(image_generation, "emit_iterm_image") as emit_image,
+    ):
+        assert (
+            register_callbacks._handle_custom_command(
+                "/codex-imagegen a red fox", "codex-imagegen"
+            )
+            is True
+        )
+
+    generate.assert_called_once_with("a red fox")
+    success.assert_called_once()
+    emit_image.assert_called_once_with(Path("/tmp/generated.png"))
+
+
+def test_imagegen_skill_and_tool_registration():
+    from code_puppy_core_plugins.chatgpt_oauth import register_callbacks
+
+    with patch.object(
+        register_callbacks,
+        "load_stored_tokens",
+        return_value={"access_token": "token", "account_id": "account"},
+    ):
+        skills = register_callbacks._register_imagegen_skill()
+        assert skills[0]["name"] == "codex-imagegen"
+        assert Path(skills[0]["skill_md_path"]).is_file()
+        assert register_callbacks._advertise_imagegen_tool("code-puppy") == [
+            "codex_imagegen"
+        ]
+    tools = register_callbacks._register_imagegen_tools()
+    assert tools == [
+        {"name": "codex_imagegen", "register_func": image_tool.register_codex_imagegen}
+    ]
+
+
+def test_imagegen_skill_is_not_registered_when_logged_out():
+    from code_puppy_core_plugins.chatgpt_oauth import register_callbacks
+
+    with patch.object(register_callbacks, "load_stored_tokens", return_value=None):
+        assert register_callbacks._register_imagegen_skill() == []
+
+
+def test_imagegen_tool_is_not_advertised_when_logged_out():
+    from code_puppy_core_plugins.chatgpt_oauth import register_callbacks
+
+    with patch.object(register_callbacks, "load_stored_tokens", return_value=None):
+        assert register_callbacks._advertise_imagegen_tool("code-puppy") == []
+
+
+def test_logout_reloads_agent_to_unbind_imagegen(tmp_path):
+    from code_puppy_core_plugins.chatgpt_oauth import register_callbacks
+
+    token_path = tmp_path / "tokens.json"
+    token_path.write_text("{}")
+    with (
+        patch.object(
+            register_callbacks,
+            "load_stored_tokens",
+            return_value={"access_token": "token", "account_id": "account"},
+        ),
+        patch.object(
+            register_callbacks, "get_token_storage_path", return_value=token_path
+        ),
+        patch.object(register_callbacks, "remove_chatgpt_models", return_value=0),
+        patch.object(register_callbacks, "emit_info"),
+        patch.object(register_callbacks, "emit_success"),
+        patch.object(register_callbacks, "_reload_active_agent") as reload_agent,
+        patch.object(register_callbacks, "refresh_skill_cache") as refresh_skills,
+    ):
+        register_callbacks._handle_chatgpt_logout()
+
+    assert not token_path.exists()
+    reload_agent.assert_called_once_with()
+    refresh_skills.assert_called_once_with()
+
+
+def test_logout_does_not_refresh_skills_when_already_logged_out(tmp_path):
+    from code_puppy_core_plugins.chatgpt_oauth import register_callbacks
+
+    token_path = tmp_path / "tokens.json"
+    with (
+        patch.object(register_callbacks, "load_stored_tokens", return_value=None),
+        patch.object(
+            register_callbacks, "get_token_storage_path", return_value=token_path
+        ),
+        patch.object(register_callbacks, "remove_chatgpt_models", return_value=0),
+        patch.object(register_callbacks, "emit_info"),
+        patch.object(register_callbacks, "emit_success"),
+        patch.object(register_callbacks, "_reload_active_agent") as reload_agent,
+        patch.object(register_callbacks, "refresh_skill_cache") as refresh_skills,
+    ):
+        register_callbacks._handle_chatgpt_logout()
+
+    reload_agent.assert_not_called()
+    refresh_skills.assert_not_called()
+
+
+def test_auth_command_refreshes_skill_cache():
+    from code_puppy_core_plugins.chatgpt_oauth import register_callbacks
+
+    with (
+        patch.object(register_callbacks, "run_oauth_flow") as oauth_flow,
+        patch.object(register_callbacks, "set_model_and_reload_agent") as set_model,
+        patch.object(register_callbacks, "refresh_skill_cache") as refresh_skills,
+    ):
+        assert (
+            register_callbacks._handle_custom_command("/codex-auth", "codex-auth")
+            is True
+        )
+
+    oauth_flow.assert_called_once_with()
+    set_model.assert_called_once_with("codex-gpt-5.6-sol")
+    refresh_skills.assert_called_once_with()
+
+
+def test_codex_imagegen_agent_tool(tmp_path):
+    registered = {}
+
+    class FakeAgent:
+        def tool(self, function):
+            registered[function.__name__] = function
+            return function
+
+    image_tool.register_codex_imagegen(FakeAgent())
+    output_path = tmp_path / "generated.png"
+    with (
+        patch.object(
+            image_tool, "generate_image", return_value=output_path
+        ) as generate,
+        patch.object(image_tool, "emit_iterm_image", return_value=True),
+    ):
+        result = asyncio.run(registered["codex_imagegen"](MagicMock(), "a fox"))
+
+    assert result == {
+        "success": True,
+        "path": str(output_path),
+        "displayed_inline": True,
+    }
+    generate.assert_called_once_with("a fox", None)
+
+
+def test_codex_imagegen_agent_tool_forwards_reference_images(tmp_path):
+    registered = {}
+
+    class FakeAgent:
+        def tool(self, function):
+            registered[function.__name__] = function
+            return function
+
+    image_tool.register_codex_imagegen(FakeAgent())
+    output_path = tmp_path / "generated.png"
+    refs = [str(tmp_path / "ref.png")]
+    with (
+        patch.object(
+            image_tool, "generate_image", return_value=output_path
+        ) as generate,
+        patch.object(image_tool, "emit_iterm_image", return_value=False),
+    ):
+        result = asyncio.run(
+            registered["codex_imagegen"](MagicMock(), "same fox, at night", refs)
+        )
+
+    assert result["success"] is True
+    generate.assert_called_once_with("same fox, at night", refs)
+
+
+def test_config_paths():
+    """Test configuration path helpers."""
+    token_path = config.get_token_storage_path()
+    assert token_path.name == "chatgpt_oauth.json"
+    # XDG paths use "code_puppy" (without dot) in ~/.local/share or ~/.config
+    assert "code_puppy" in str(token_path)
+
+    config_dir = config.get_config_dir()
+    # Default is ~/.code_puppy; XDG paths only used when XDG env vars are set
+    assert config_dir.name in ("code_puppy", ".code_puppy")
+
+    chatgpt_models = config.get_chatgpt_models_path()
+    assert chatgpt_models.name == "chatgpt_models.json"
+
+
+def test_oauth_config():
+    """Test OAuth configuration values."""
+    assert config.CHATGPT_OAUTH_CONFIG["issuer"] == "https://auth.openai.com"
+    assert config.CHATGPT_OAUTH_CONFIG["client_id"] == "app_EMoamEEZ73f0CkXaXp7hrann"
+    assert config.CHATGPT_OAUTH_CONFIG["prefix"] == "codex-"
+
+
+def test_jwt_parsing_with_nested_org():
+    """Test JWT parsing with nested organization structure like the user's payload."""
+    # This simulates the user's JWT payload structure
+    mock_claims = {
+        "aud": ["app_EMoamEEZ73f0CkXaXp7hrann"],
+        "auth_provider": "google",
+        "email": "mike.pfaf fenberger@gmail.com",
+        "https://api.openai.com/auth": {
+            "chatgpt_account_id": "d1844a91-9aac-419b-903e-f6a99c76f163",
+            "organizations": [
+                {
+                    "id": "org-iydWjnSxSr51VuYhDVMDte5",
+                    "is_default": True,
+                    "role": "owner",
+                    "title": "Personal",
+                }
+            ],
+            "groups": ["api-data-sharing-incentives-program", "verified-organization"],
+        },
+        "sub": "google-oauth2|107692466937587138174",
+    }
+
+    # Test the org extraction logic
+    auth_claims = mock_claims.get("https://api.openai.com/auth", {})
+    organizations = auth_claims.get("organizations", [])
+
+    org_id = None
+    if organizations:
+        default_org = next(
+            (org for org in organizations if org.get("is_default")), organizations[0]
+        )
+        org_id = default_org.get("id")
+
+    assert org_id == "org-iydWjnSxSr51VuYhDVMDte5"
+
+    # Test fallback to top-level org_id (should not happen in this case)
+    if not org_id:
+        org_id = mock_claims.get("organization_id")
+
+    assert org_id == "org-iydWjnSxSr51VuYhDVMDte5"
+    assert config.CHATGPT_OAUTH_CONFIG["required_port"] == 1455
+
+
+def test_code_verifier_generation():
+    """Test PKCE code verifier generation."""
+    verifier = utils._generate_code_verifier()
+    assert isinstance(verifier, str)
+    assert len(verifier) > 50  # Should be long
+
+
+def test_code_challenge_computation():
+    """Test PKCE code challenge computation."""
+    verifier = "test_verifier_string"
+    challenge = utils._compute_code_challenge(verifier)
+    assert isinstance(challenge, str)
+    assert len(challenge) > 0
+    # Should be URL-safe base64
+    assert all(
+        c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
+        for c in challenge
+    )
+
+
+def test_prepare_oauth_context():
+    """Test OAuth context preparation."""
+    context = utils.prepare_oauth_context()
+    assert context.state
+    assert context.code_verifier
+    assert context.code_challenge
+    assert context.created_at > 0
+    assert context.redirect_uri is None
+
+
+def test_assign_redirect_uri():
+    """Test redirect URI assignment."""
+    context = utils.prepare_oauth_context()
+    redirect_uri = utils.assign_redirect_uri(context, 1455)
+    assert redirect_uri == "http://localhost:1455/auth/callback"
+    assert context.redirect_uri == redirect_uri
+
+
+def test_build_authorization_url():
+    """Test authorization URL building."""
+    context = utils.prepare_oauth_context()
+    utils.assign_redirect_uri(context, 1455)
+    auth_url = utils.build_authorization_url(context)
+
+    assert auth_url.startswith("https://auth.openai.com/oauth/authorize?")
+    assert "response_type=code" in auth_url
+    assert "client_id=" in auth_url
+    assert "redirect_uri=" in auth_url
+    assert "code_challenge=" in auth_url
+    assert "code_challenge_method=S256" in auth_url
+    assert f"state={context.state}" in auth_url
+
+
+def test_parse_jwt_claims():
+    """Test JWT claims parsing."""
+    # Valid JWT structure (header.payload.signature)
+    import base64
+
+    payload = base64.urlsafe_b64encode(json.dumps({"sub": "user123"}).encode()).decode()
+    token = f"header.{payload}.signature"
+
+    claims = utils.parse_jwt_claims(token)
+    assert claims is not None
+    assert claims["sub"] == "user123"
+
+    # Invalid token
+    assert utils.parse_jwt_claims("") is None
+    assert utils.parse_jwt_claims("invalid") is None
+
+
+def test_save_and_load_tokens(tmp_path):
+    """Test token storage and retrieval."""
+    with patch.object(
+        utils, "get_token_storage_path", return_value=tmp_path / "tokens.json"
+    ):
+        tokens = {
+            "access_token": "test_access",
+            "refresh_token": "test_refresh",
+            "api_key": "sk-test",
+        }
+
+        # Save tokens
+        assert utils.save_tokens(tokens)
+
+        # Load tokens
+        loaded = utils.load_stored_tokens()
+        assert loaded == tokens
+
+
+def test_save_and_load_chatgpt_models(tmp_path):
+    """Test ChatGPT models configuration."""
+    with patch.object(
+        utils, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json"
+    ):
+        models = {
+            "codex-gpt-4o": {
+                "type": "openai",
+                "name": "gpt-4o",
+                "oauth_source": "chatgpt-oauth-plugin",
+            }
+        }
+
+        # Save models
+        assert utils.save_chatgpt_models(models)
+
+        # Load models
+        loaded = utils.load_chatgpt_models()
+        assert loaded == models
+
+
+def test_remove_chatgpt_models(tmp_path):
+    """Test removal of ChatGPT models from config."""
+    with patch.object(
+        utils, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json"
+    ):
+        models = {
+            "codex-gpt-4o": {
+                "type": "openai",
+                "oauth_source": "chatgpt-oauth-plugin",
+            },
+            "claude-3-opus": {
+                "type": "anthropic",
+                "oauth_source": "other",
+            },
+        }
+        utils.save_chatgpt_models(models)
+
+        # Remove only ChatGPT models
+        removed_count = utils.remove_chatgpt_models()
+        assert removed_count == 1
+
+        # Verify only ChatGPT model was removed
+        remaining = utils.load_chatgpt_models()
+        assert "codex-gpt-4o" not in remaining
+        assert "claude-3-opus" in remaining
+
+
+@patch("code_puppy_core_plugins.chatgpt_oauth.utils.requests.post")
+def test_exchange_code_for_tokens(mock_post):
+    """Test authorization code exchange."""
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {
+        "access_token": "test_access",
+        "refresh_token": "test_refresh",
+        "id_token": "test_id",
+    }
+    mock_post.return_value = mock_response
+
+    context = utils.prepare_oauth_context()
+    utils.assign_redirect_uri(context, 1455)
+
+    tokens = utils.exchange_code_for_tokens("test_code", context)
+    assert tokens is not None
+    assert tokens["access_token"] == "test_access"
+    assert "last_refresh" in tokens
+
+
+@patch("code_puppy_core_plugins.chatgpt_oauth.utils.requests.get")
+def test_fetch_chatgpt_models(mock_get):
+    """Test fetching models from ChatGPT Codex API."""
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    # New response format uses "models" key with "slug" field
+    mock_response.json.return_value = {
+        "models": [
+            {"slug": "gpt-4o"},
+            {"slug": "gpt-3.5-turbo"},
+            {"slug": "o1-preview"},
+            {"slug": "codex-mini"},
+        ]
+    }
+    mock_get.return_value = mock_response
+
+    models = utils.fetch_chatgpt_models("test_access_token", "test_account_id")
+    assert models is not None
+    # A successful endpoint response is authoritative; fallback models are
+    # used only when discovery fails.
+    assert "gpt-5.4" not in models
+    assert "gpt-5.4-mini" not in models
+    # API-returned models are preserved.
+    assert "gpt-4o" in models
+    assert "gpt-3.5-turbo" in models
+    assert "o1-preview" in models
+    assert "codex-mini" in models
+
+
+@patch("code_puppy_core_plugins.chatgpt_oauth.utils.requests.get")
+def test_fetch_chatgpt_models_fallback(mock_get):
+    """Test that fetch_chatgpt_models returns default list on API failure."""
+    mock_response = MagicMock()
+    mock_response.status_code = 404
+    mock_response.text = '{"detail":"Not Found"}'
+    mock_get.return_value = mock_response
+
+    models = utils.fetch_chatgpt_models("test_access_token", "test_account_id")
+    assert models is not None
+    # Keep the fallback aligned with the models currently exposed by Codex.
+    assert "gpt-5.4" in models
+    assert "gpt-5.4-mini" in models
+    assert "gpt-5.3-codex-spark" in models
+    assert "codex-auto-review" in models
+    assert "gpt-5.3-instant" not in models
+    assert "gpt-5.3-codex" not in models
+    assert "gpt-5.2-codex" not in models
+    assert "gpt-5.2" not in models
+
+
+def test_add_models_to_chatgpt_config(tmp_path):
+    """Test adding models to chatgpt_models.json."""
+    with patch.object(
+        utils, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json"
+    ):
+        models = ["gpt-4o", "gpt-3.5-turbo"]
+
+        assert utils.add_models_to_extra_config(models)
+
+        loaded = utils.load_chatgpt_models()
+        assert "codex-gpt-4o" in loaded
+        assert "codex-gpt-3.5-turbo" in loaded
+        assert loaded["codex-gpt-4o"]["type"] == "chatgpt_oauth"
+        assert loaded["codex-gpt-4o"]["name"] == "gpt-4o"
+        assert loaded["codex-gpt-4o"]["oauth_source"] == "chatgpt-oauth-plugin"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])