From 77aa3b85d15c708989baaae52ccdfc28e3a0d617 Mon Sep 17 00:00:00 2001 From: amithkoujalgi Date: Sat, 30 Aug 2025 15:07:32 +0530 Subject: [PATCH] Refactor WithAuth integration test for improved structure and functionality - Restored the WithAuth class and updated it to include a new general-purpose model. - Enhanced the setup method to initialize the OllamaAPI with bearer authentication. - Added tests to validate the API's behavior with correct and incorrect authentication tokens. - Updated the structured output test to reflect changes in the model and prompt. - Improved logging for better traceability during test execution. --- .../ollama4j/integrationtests/WithAuth.java | 374 +++++++++--------- 1 file changed, 180 insertions(+), 194 deletions(-) diff --git a/src/test/java/io/github/ollama4j/integrationtests/WithAuth.java b/src/test/java/io/github/ollama4j/integrationtests/WithAuth.java index f14e592..4b3d0c7 100644 --- a/src/test/java/io/github/ollama4j/integrationtests/WithAuth.java +++ b/src/test/java/io/github/ollama4j/integrationtests/WithAuth.java @@ -1,194 +1,180 @@ -//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; -//import org.junit.jupiter.api.MethodOrderer.OrderAnnotation; -//import org.junit.jupiter.api.Order; -//import org.junit.jupiter.api.Test; -//import org.junit.jupiter.api.TestMethodOrder; -//import org.slf4j.Logger; -//import org.slf4j.LoggerFactory; -//import org.testcontainers.containers.GenericContainer; -//import org.testcontainers.containers.NginxContainer; -//import org.testcontainers.containers.wait.strategy.Wait; -//import org.testcontainers.ollama.OllamaContainer; -//import org.testcontainers.utility.DockerImageName; -//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.*; -// -//@OllamaToolService(providers = {AnnotatedTool.class}) -//@TestMethodOrder(OrderAnnotation.class) -//@SuppressWarnings({"HttpUrlsUsage", "SpellCheckingInspection", "resource", "ResultOfMethodCallIgnored"}) -//public class WithAuth { -// -// private static final Logger LOG = LoggerFactory.getLogger(WithAuth.class); -// 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; -// private static OllamaAPI api; -// -// @BeforeAll -// public static void setUp() { -// ollama = createOllamaContainer(); -// ollama.start(); -// -// nginx = createNginxContainer(ollama.getMappedPort(OLLAMA_INTERNAL_PORT)); -// nginx.start(); -// -// LOG.info("Using Testcontainer Ollama host..."); -// -// api = new OllamaAPI("http://" + nginx.getHost() + ":" + nginx.getMappedPort(NGINX_PORT)); -// 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() { -// return new OllamaContainer("ollama/ollama:" + OLLAMA_VERSION).withExposedPorts(OLLAMA_INTERNAL_PORT); -// } -// -// private static String generateNginxConfig(int ollamaPort) { -// return String.format("events {}\n" + -// "\n" + -// "http {\n" + -// " server {\n" + -// " listen 80;\n" + -// "\n" + -// " location / {\n" + -// " set $auth_header $http_authorization;\n" + -// "\n" + -// " if ($auth_header != \"Bearer secret-token\") {\n" + -// " return 401;\n" + -// " }\n" + -// "\n" + -// " proxy_pass http://host.docker.internal:%s/;\n" + -// " proxy_set_header Host $host;\n" + -// " proxy_set_header X-Real-IP $remote_addr;\n" + -// " proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;\n" + -// " proxy_set_header X-Forwarded-Proto $scheme;\n" + -// " }\n" + -// " }\n" + -// "}\n", ollamaPort); -// } -// -// 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(); -// -// nginxConf = new File(tempDir, "nginx.conf"); -// try (FileWriter writer = new FileWriter(nginxConf)) { -// writer.write(generateNginxConfig(ollamaPort)); -// } -// -// return new NginxContainer<>(DockerImageName.parse(NGINX_VERSION)) -// .withExposedPorts(NGINX_PORT) -// .withCopyFileToContainer( -// MountableFile.forHostPath(nginxConf.getAbsolutePath()), -// "/etc/nginx/nginx.conf" -// ) -// .withExtraHost("host.docker.internal", "host-gateway") -// .waitingFor( -// Wait.forHttp("/") -// .forStatusCode(401) -// .withStartupTimeout(Duration.ofSeconds(30)) -// ); -// } catch (IOException e) { -// throw new RuntimeException("Failed to create nginx.conf", e); -// } -// } -// -// @Test -// @Order(1) -// 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 format = new HashMap<>(); -// format.put("type", "object"); -// format.put("properties", new HashMap() { -// { -// put("timeHour", new HashMap() { -// { -// put("type", "integer"); -// } -// }); -// put("isNightTime", new HashMap() { -// { -// 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()); -// } -//} +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; +import org.junit.jupiter.api.MethodOrderer.OrderAnnotation; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestMethodOrder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.NginxContainer; +import org.testcontainers.containers.wait.strategy.Wait; +import org.testcontainers.ollama.OllamaContainer; +import org.testcontainers.utility.DockerImageName; +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.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.*; + +@OllamaToolService(providers = {AnnotatedTool.class}) +@TestMethodOrder(OrderAnnotation.class) +@SuppressWarnings({"HttpUrlsUsage", "SpellCheckingInspection", "resource", "ResultOfMethodCallIgnored"}) +public class WithAuth { + + private static final Logger LOG = LoggerFactory.getLogger(WithAuth.class); + 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 GENERAL_PURPOSE_MODEL = "gemma3:270m"; + + + private static OllamaContainer ollama; + private static GenericContainer nginx; + private static OllamaAPI api; + + @BeforeAll + public static void setUp() { + ollama = createOllamaContainer(); + ollama.start(); + + nginx = createNginxContainer(ollama.getMappedPort(OLLAMA_INTERNAL_PORT)); + nginx.start(); + + LOG.info("Using Testcontainer Ollama host..."); + + api = new OllamaAPI("http://" + nginx.getHost() + ":" + nginx.getMappedPort(NGINX_PORT)); + 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() { + return new OllamaContainer("ollama/ollama:" + OLLAMA_VERSION).withExposedPorts(OLLAMA_INTERNAL_PORT); + } + + private static String generateNginxConfig(int ollamaPort) { + return String.format("events {}\n" + + "\n" + + "http {\n" + + " server {\n" + + " listen 80;\n" + + "\n" + + " location / {\n" + + " set $auth_header $http_authorization;\n" + + "\n" + + " if ($auth_header != \"Bearer secret-token\") {\n" + + " return 401;\n" + + " }\n" + + "\n" + + " proxy_pass http://host.docker.internal:%s/;\n" + + " proxy_set_header Host $host;\n" + + " proxy_set_header X-Real-IP $remote_addr;\n" + + " proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;\n" + + " proxy_set_header X-Forwarded-Proto $scheme;\n" + + " }\n" + + " }\n" + + "}\n", ollamaPort); + } + + 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(); + + nginxConf = new File(tempDir, "nginx.conf"); + try (FileWriter writer = new FileWriter(nginxConf)) { + writer.write(generateNginxConfig(ollamaPort)); + } + + return new NginxContainer<>(DockerImageName.parse(NGINX_VERSION)) + .withExposedPorts(NGINX_PORT) + .withCopyFileToContainer( + MountableFile.forHostPath(nginxConf.getAbsolutePath()), + "/etc/nginx/nginx.conf" + ) + .withExtraHost("host.docker.internal", "host-gateway") + .waitingFor( + Wait.forHttp("/") + .forStatusCode(401) + .withStartupTimeout(Duration.ofSeconds(30)) + ); + } catch (IOException e) { + throw new RuntimeException("Failed to create nginx.conf", e); + } + } + + @Test + @Order(1) + 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(GENERAL_PURPOSE_MODEL); + + String prompt = "The sun is shining brightly and is directly overhead at the zenith, casting my shadow over my foot, so it must be noon."; + + Map format = new HashMap<>(); + format.put("type", "object"); + format.put("properties", new HashMap() { + { + put("isNoon", new HashMap() { + { + put("type", "boolean"); + } + }); + } + }); + format.put("required", List.of("isNoon")); + + OllamaResult result = api.generate(GENERAL_PURPOSE_MODEL, prompt, format); + + assertNotNull(result); + assertNotNull(result.getResponse()); + assertFalse(result.getResponse().isEmpty()); + + assertEquals(true, result.getStructuredResponse().get("isNoon")); + } +}