diff --git a/src/test/java/io/github/ollama4j/integrationtests/OllamaAPIIntegrationTest.java b/src/test/java/io/github/ollama4j/integrationtests/OllamaAPIIntegrationTest.java
index e19ed0a..1973912 100644
--- a/src/test/java/io/github/ollama4j/integrationtests/OllamaAPIIntegrationTest.java
+++ b/src/test/java/io/github/ollama4j/integrationtests/OllamaAPIIntegrationTest.java
@@ -395,10 +395,7 @@ class OllamaAPIIntegrationTest {
                 chatResult.getResponseModel().getMessage().getRole().getRoleName(),
                 "Role of the response message should be ASSISTANT");
         List<OllamaChatToolCalls> toolCalls = chatResult.getChatHistory().get(1).getToolCalls();
-        assertEquals(
-                1,
-                toolCalls.size(),
-                "There should be exactly one tool call in the second chat history message");
+        assert (!toolCalls.isEmpty());
         OllamaToolCallsFunction function = toolCalls.get(0).getFunction();
         assertEquals(
                 "get-employee-details",
@@ -979,7 +976,10 @@ class OllamaAPIIntegrationTest {
                                 LOG.info(
                                         "Invoking employee finder tool with arguments: {}",
                                         arguments);
-                                String employeeName = arguments.get("employee-name").toString();
+                                String employeeName = "Random Employee";
+                                if (arguments.containsKey("employee-name")) {
+                                    employeeName = arguments.get("employee-name").toString();
+                                }
                                 String address = null;
                                 String phone = null;
                                 if (employeeName.equalsIgnoreCase("Rahul Kumar")) {