Skip to content

Commit f0f168e

Browse files
committed
mocked unit tests for providers
1 parent 3ba3efd commit f0f168e

2 files changed

Lines changed: 276 additions & 3 deletions

File tree

ballerina-interpreter/tests/agent_test.bal

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import ballerina/http;
1818
import ballerina/io;
1919
import ballerina/lang.runtime;
2020
import ballerina/test;
21+
import ballerina/ai;
2122

2223
@test:Config
2324
function testValidateJsonSchemaNullSchema() returns error? {
@@ -358,9 +359,7 @@ function testExtractJsonFromCodeBlock(string description, string response, strin
358359
test:assertEquals(result, expected);
359360
}
360361

361-
// ============================================
362-
// Array Schema End-to-End Tests
363-
// ============================================
362+
364363

365364
@test:Config
366365
function testArrayOutputSchemaEndToEnd() returns error? {
@@ -410,3 +409,24 @@ function testArrayOutputSchemaInvalidResponse() returns error? {
410409
// Should return 500 error due to schema validation failure
411410
test:assertEquals(response.statusCode, 500, "Should return 500 for schema validation failure");
412411
}
412+
413+
@test:Config
414+
function testGetModelOllamaLocalLoopback() returns error? {
415+
Model model = {
416+
provider: "ollama",
417+
name: "llama3"
418+
};
419+
var result = getModel(model);
420+
test:assertTrue(result is ai:ModelProvider);
421+
}
422+
423+
@test:Config
424+
function testGetModelOllamaCustomEndpoint() returns error? {
425+
Model model = {
426+
provider: "ollama",
427+
name: "mistral",
428+
url: "http://192.168.1.15:11434"
429+
};
430+
var result = getModel(model);
431+
test:assertTrue(result is ai:ModelProvider);
432+
}
Lines changed: 253 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,253 @@
1+
# Copyright (c) 2026, WSO2 LLC. (https://www.wso2.com).
2+
#
3+
# WSO2 LLC. licenses this file to you under the Apache License,
4+
# Version 2.0 (the "License"); you may not use this file except
5+
# in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing,
11+
# software distributed under the License is distributed on an
12+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
13+
# KIND, either express or implied. See the License for the
14+
# specific language governing permissions and limitations
15+
# under the License.
16+
17+
import os
18+
import sys
19+
from unittest.mock import MagicMock, patch
20+
21+
import pytest
22+
23+
from afm.exceptions import ProviderError
24+
from afm.models import ClientAuthentication, Model
25+
from afm_langchain.providers import create_model_provider
26+
27+
28+
class TestOpenAIProvider:
29+
@patch("langchain_openai.ChatOpenAI")
30+
def test_create_openai_with_explicit_credentials(self, mock_chat_openai: MagicMock) -> None:
31+
"""Verify OpenAI client creation with credentials passed in AFM metadata."""
32+
afm_model = Model(
33+
provider="openai",
34+
name="gpt-4o",
35+
authentication=ClientAuthentication(
36+
type="api-key",
37+
api_key="explicit-openai-token-123"
38+
)
39+
)
40+
create_model_provider(afm_model)
41+
42+
mock_chat_openai.assert_called_once()
43+
kwargs = mock_chat_openai.call_args[1]
44+
assert kwargs["model"] == "gpt-4o"
45+
assert kwargs["api_key"] == "explicit-openai-token-123"
46+
47+
@patch("langchain_openai.ChatOpenAI")
48+
def test_create_openai_with_env_fallback(self, mock_chat_openai: MagicMock) -> None:
49+
"""Verify OpenAI client successfully falls back to environment variables."""
50+
afm_model = Model(provider="openai", name="gpt-4o")
51+
52+
with patch.dict(os.environ, {"OPENAI_API_KEY": "env-openai-token-456"}):
53+
create_model_provider(afm_model)
54+
55+
mock_chat_openai.assert_called_once()
56+
kwargs = mock_chat_openai.call_args[1]
57+
assert kwargs["api_key"] == "env-openai-token-456"
58+
59+
def test_create_openai_missing_credentials_fails(self) -> None:
60+
"""Verify ProviderError is raised when no API key can be resolved."""
61+
afm_model = Model(provider="openai", name="gpt-4o")
62+
63+
with patch.dict(os.environ, {}, clear=True):
64+
with pytest.raises(ProviderError) as exc_info:
65+
create_model_provider(afm_model)
66+
assert "No API key found" in str(exc_info.value)
67+
assert exc_info.value.provider == "openai"
68+
69+
def test_openai_import_error(self) -> None:
70+
"""Verify clean package warning when the underlying langchain library is missing."""
71+
afm_model = Model(
72+
provider="openai",
73+
name="gpt-4o",
74+
authentication=ClientAuthentication(type="api-key", api_key="dummy")
75+
)
76+
with patch.dict(sys.modules, {"langchain_openai": None}):
77+
with pytest.raises(ProviderError) as exc_info:
78+
create_model_provider(afm_model)
79+
assert "langchain-openai package is required" in str(exc_info.value)
80+
81+
82+
class TestAnthropicProvider:
83+
@patch("langchain_anthropic.ChatAnthropic")
84+
def test_create_anthropic_with_explicit_credentials(self, mock_chat_anthropic: MagicMock) -> None:
85+
"""Verify Anthropic client initialization with explicit credentials."""
86+
afm_model = Model(
87+
provider="anthropic",
88+
name="claude-sonnet-4-5",
89+
authentication=ClientAuthentication(
90+
type="api-key",
91+
api_key="explicit-anthropic-token-123"
92+
)
93+
)
94+
create_model_provider(afm_model)
95+
96+
mock_chat_anthropic.assert_called_once()
97+
kwargs = mock_chat_anthropic.call_args[1]
98+
assert kwargs["model"] == "claude-sonnet-4-5"
99+
assert kwargs["api_key"] == "explicit-anthropic-token-123"
100+
101+
@patch("langchain_anthropic.ChatAnthropic")
102+
def test_create_anthropic_with_env_fallback(self, mock_chat_anthropic: MagicMock) -> None:
103+
"""Verify Anthropic client falls back onto system environment variables."""
104+
afm_model = Model(provider="anthropic", name="claude-sonnet-4-5")
105+
106+
with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "env-anthropic-token-456"}):
107+
create_model_provider(afm_model)
108+
109+
mock_chat_anthropic.assert_called_once()
110+
kwargs = mock_chat_anthropic.call_args[1]
111+
assert kwargs["api_key"] == "env-anthropic-token-456"
112+
113+
def test_anthropic_import_error(self) -> None:
114+
"""Verify clean package warning when langchain-anthropic is missing."""
115+
afm_model = Model(
116+
provider="anthropic",
117+
name="claude-sonnet-4-5",
118+
authentication=ClientAuthentication(type="api-key", api_key="dummy")
119+
)
120+
with patch.dict(sys.modules, {"langchain_anthropic": None}):
121+
with pytest.raises(ProviderError) as exc_info:
122+
create_model_provider(afm_model)
123+
assert "langchain-anthropic package is required" in str(exc_info.value)
124+
125+
126+
class TestGeminiProvider:
127+
@patch("langchain_google_genai.ChatGoogleGenerativeAI")
128+
def test_create_gemini_studio_explicit_credentials(self, mock_chat_gemini: MagicMock) -> None:
129+
"""Verify standard AI Studio pathway accepts and passes API keys cleanly."""
130+
afm_model = Model(
131+
provider="gemini",
132+
name="gemini-2.5-flash",
133+
authentication=ClientAuthentication(
134+
type="api-key",
135+
api_key="explicit-gemini-token-789"
136+
)
137+
)
138+
create_model_provider(afm_model)
139+
140+
mock_chat_gemini.assert_called_once()
141+
kwargs = mock_chat_gemini.call_args[1]
142+
assert kwargs["model"] == "gemini-2.5-flash"
143+
assert kwargs["api_key"] == "explicit-gemini-token-789"
144+
assert "vertexai" not in kwargs
145+
146+
@patch("langchain_google_genai.ChatGoogleGenerativeAI")
147+
def test_create_gemini_studio_env_fallback(self, mock_chat_gemini: MagicMock) -> None:
148+
"""Verify standard AI Studio pathway falls back onto GOOGLE_API_KEY environment variables."""
149+
afm_model = Model(provider="gemini", name="gemini-2.5-flash")
150+
151+
with patch.dict(os.environ, {"GOOGLE_API_KEY": "env-gemini-token-101"}):
152+
create_model_provider(afm_model)
153+
154+
mock_chat_gemini.assert_called_once()
155+
kwargs = mock_chat_gemini.call_args[1]
156+
assert kwargs["api_key"] == "env-gemini-token-101"
157+
158+
@patch("langchain_google_genai.ChatGoogleGenerativeAI")
159+
def test_create_gemini_vertex_adc_fallback(self, mock_chat_gemini: MagicMock) -> None:
160+
"""Verify enterprise Vertex AI path bypasses API key lookups to allow Application Default Credentials (ADC)."""
161+
afm_model = Model(
162+
provider="gemini",
163+
name="gemini-2.5-flash",
164+
project="wso2-gcp-enterprise-sandbox",
165+
location="us-central1"
166+
)
167+
create_model_provider(afm_model)
168+
169+
mock_chat_gemini.assert_called_once()
170+
kwargs = mock_chat_gemini.call_args[1]
171+
assert kwargs["project"] == "wso2-gcp-enterprise-sandbox"
172+
assert kwargs["location"] == "us-central1"
173+
assert kwargs["vertexai"] is True
174+
assert "api_key" not in kwargs
175+
176+
@patch("langchain_google_genai.ChatGoogleGenerativeAI")
177+
def test_create_gemini_vertex_with_explicit_credentials(self, mock_chat_gemini: MagicMock) -> None:
178+
"""Verify enterprise Vertex AI pathways still accept explicit credentials if provided."""
179+
afm_model = Model(
180+
provider="gemini",
181+
name="gemini-2.5-flash",
182+
project="wso2-gcp-enterprise-sandbox",
183+
location="us-central1",
184+
authentication=ClientAuthentication(
185+
type="api-key",
186+
api_key="vertex-specific-explicit-key"
187+
)
188+
)
189+
create_model_provider(afm_model)
190+
191+
mock_chat_gemini.assert_called_once()
192+
kwargs = mock_chat_gemini.call_args[1]
193+
assert kwargs["project"] == "wso2-gcp-enterprise-sandbox"
194+
assert kwargs["vertexai"] is True
195+
assert kwargs["api_key"] == "vertex-specific-explicit-key"
196+
197+
def test_gemini_import_error(self) -> None:
198+
"""Verify clean package warning when langchain-google-genai is missing."""
199+
afm_model = Model(
200+
provider="gemini",
201+
name="gemini-2.5-flash",
202+
authentication=ClientAuthentication(type="api-key", api_key="dummy")
203+
)
204+
with patch.dict(sys.modules, {"langchain_google_genai": None}):
205+
with pytest.raises(ProviderError) as exc_info:
206+
create_model_provider(afm_model)
207+
assert "langchain-google-genai package is required" in str(exc_info.value)
208+
209+
210+
class TestOllamaProvider:
211+
@patch("langchain_ollama.ChatOllama")
212+
def test_create_ollama_local_loopback_default(self, mock_chat_ollama: MagicMock) -> None:
213+
"""Verify Ollama initialization functions completely unauthenticated with loopback defaults."""
214+
afm_model = Model(provider="ollama", name="llama3")
215+
create_model_provider(afm_model)
216+
217+
mock_chat_ollama.assert_called_once()
218+
kwargs = mock_chat_ollama.call_args[1]
219+
assert kwargs["model"] == "llama3"
220+
assert kwargs["base_url"] == "http://localhost:11434"
221+
assert "api_key" not in kwargs
222+
223+
@patch("langchain_ollama.ChatOllama")
224+
def test_create_ollama_with_custom_endpoint(self, mock_chat_ollama: MagicMock) -> None:
225+
"""Verify Ollama correctly respects and binds custom URL parameters (useful for container network links)."""
226+
afm_model = Model(
227+
provider="ollama",
228+
name="mistral",
229+
url="http://host.docker.internal:11434"
230+
)
231+
create_model_provider(afm_model)
232+
233+
mock_chat_ollama.assert_called_once()
234+
kwargs = mock_chat_ollama.call_args[1]
235+
assert kwargs["model"] == "mistral"
236+
assert kwargs["base_url"] == "http://host.docker.internal:11434"
237+
238+
def test_ollama_import_error(self) -> None:
239+
"""Verify clean package warning when langchain-ollama is missing."""
240+
afm_model = Model(provider="ollama", name="llama3")
241+
with patch.dict(sys.modules, {"langchain_ollama": None}):
242+
with pytest.raises(ProviderError) as exc_info:
243+
create_model_provider(afm_model)
244+
assert "langchain-ollama package is required" in str(exc_info.value)
245+
246+
247+
class TestGeneralProviderErrors:
248+
def test_unsupported_provider_raises_error(self) -> None:
249+
"""Verify ProviderError triggers upon encountering an unknown provider name."""
250+
afm_model = Model(provider="unsupported-model-ecosystem", name="test-model")
251+
with pytest.raises(ProviderError) as exc_info:
252+
create_model_provider(afm_model)
253+
assert "Unsupported provider: unsupported-model-ecosystem" in str(exc_info.value)

0 commit comments

Comments
 (0)