Update WithAuth.java

This commit is contained in:
amithkoujalgi 2025-08-16 09:27:40 +05:30
parent 54d8cf4cd9
commit cc13a580fa
No known key found for this signature in database
GPG Key ID: E29A37746AF94B70

View File

@ -1,6 +1,8 @@
package io.github.ollama4j.integrationtests;
import io.github.ollama4j.OllamaAPI;
import io.github.ollama4j.exceptions.OllamaBaseException;
import io.github.ollama4j.models.response.OllamaResult;
import io.github.ollama4j.samples.AnnotatedTool;
import io.github.ollama4j.tools.annotations.OllamaToolService;
import org.junit.jupiter.api.BeforeAll;
@ -20,9 +22,13 @@ import org.testcontainers.utility.MountableFile;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.net.URISyntaxException;
import java.time.Duration;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.*;
@OllamaToolService(providers = {AnnotatedTool.class})
@TestMethodOrder(OrderAnnotation.class)
@ -33,6 +39,10 @@ public class WithAuth {
private static final int NGINX_PORT = 80;
private static final int OLLAMA_INTERNAL_PORT = 11434;
private static final String OLLAMA_VERSION = "0.6.1";
private static final String NGINX_VERSION = "nginx:1.23.4-alpine";
private static final String BEARER_AUTH_TOKEN = "secret-token";
private static final String CHAT_MODEL_LLAMA3 = "llama3";
private static OllamaContainer ollama;
private static GenericContainer<?> nginx;
@ -52,12 +62,20 @@ public class WithAuth {
api.setRequestTimeoutSeconds(120);
api.setVerbose(true);
api.setNumberOfRetriesForModelPull(3);
String ollamaUrl = "http://" + ollama.getHost() + ":" + ollama.getMappedPort(OLLAMA_INTERNAL_PORT);
String nginxUrl = "http://" + nginx.getHost() + ":" + nginx.getMappedPort(NGINX_PORT);
LOG.info(
"The Ollama service is now accessible via the Nginx proxy with bearer-auth authentication mode.\n" +
"→ Ollama URL: {}\n" +
"→ Proxy URL: {}}",
ollamaUrl, nginxUrl
);
LOG.info("OllamaAPI initialized with bearer auth token: {}", BEARER_AUTH_TOKEN);
}
private static OllamaContainer createOllamaContainer() {
OllamaContainer container = new OllamaContainer("ollama/ollama:" + OLLAMA_VERSION);
container.addExposedPort(OLLAMA_INTERNAL_PORT);
return container;
return new OllamaContainer("ollama/ollama:" + OLLAMA_VERSION).withExposedPorts(OLLAMA_INTERNAL_PORT);
}
private static String generateNginxConfig(int ollamaPort) {
@ -86,7 +104,6 @@ public class WithAuth {
public static GenericContainer<?> createNginxContainer(int ollamaPort) {
File nginxConf;
try {
File tempDir = new File(System.getProperty("java.io.tmpdir"), "nginx-auth");
if (!tempDir.exists()) tempDir.mkdirs();
@ -96,7 +113,7 @@ public class WithAuth {
writer.write(generateNginxConfig(ollamaPort));
}
return new NginxContainer<>(DockerImageName.parse("nginx:1.23.4-alpine"))
return new NginxContainer<>(DockerImageName.parse(NGINX_VERSION))
.withExposedPorts(NGINX_PORT)
.withCopyFileToContainer(
MountableFile.forHostPath(nginxConf.getAbsolutePath()),
@ -115,12 +132,63 @@ public class WithAuth {
@Test
@Order(1)
void testEndpoint() throws InterruptedException {
String ollamaUrl = "http://" + ollama.getHost() + ":" + ollama.getMappedPort(OLLAMA_INTERNAL_PORT);
String nginxUrl = "http://" + nginx.getHost() + ":" + nginx.getMappedPort(NGINX_PORT);
System.out.printf("Ollama service at %s is now accessible through the Nginx proxy at %s%n", ollamaUrl, nginxUrl);
api.setBearerAuth("secret-token");
Thread.sleep(1000);
assertTrue(api.ping(), "OllamaAPI failed to ping through NGINX with auth.");
void testOllamaBehindProxy() throws InterruptedException {
api.setBearerAuth(BEARER_AUTH_TOKEN);
assertTrue(api.ping(), "Expected OllamaAPI to successfully ping through NGINX with valid auth token.");
}
@Test
@Order(1)
void testWithWrongToken() throws InterruptedException {
api.setBearerAuth("wrong-token");
assertFalse(api.ping(), "Expected OllamaAPI ping to fail through NGINX with an invalid auth token.");
}
@Test
@Order(2)
void testAskModelWithStructuredOutput()
throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
api.setBearerAuth(BEARER_AUTH_TOKEN);
api.pullModel(CHAT_MODEL_LLAMA3);
int timeHour = 6;
boolean isNightTime = false;
String prompt = "The Sun is shining, and its " + timeHour + ". Its daytime.";
Map<String, Object> format = new HashMap<>();
format.put("type", "object");
format.put("properties", new HashMap<String, Object>() {
{
put("timeHour", new HashMap<String, Object>() {
{
put("type", "integer");
}
});
put("isNightTime", new HashMap<String, Object>() {
{
put("type", "boolean");
}
});
}
});
format.put("required", Arrays.asList("timeHour", "isNightTime"));
OllamaResult result = api.generate(CHAT_MODEL_LLAMA3, prompt, format);
assertNotNull(result);
assertNotNull(result.getResponse());
assertFalse(result.getResponse().isEmpty());
assertEquals(timeHour,
result.getStructuredResponse().get("timeHour"));
assertEquals(isNightTime,
result.getStructuredResponse().get("isNightTime"));
TimeOfDay timeOfDay = result.as(TimeOfDay.class);
assertEquals(timeHour, timeOfDay.getTimeHour());
assertEquals(isNightTime, timeOfDay.isNightTime());
}
}