From d78ff1f4814bed808e20c4d4fc853c3632151a0c Mon Sep 17 00:00:00 2001
From: Amrit Singh <amritsingh374@gmail.com>
Date: Wed, 29 Jan 2025 14:56:38 +0530
Subject: [PATCH] using testify for tests and unexporting someTool

---
 tools/tool.go      |  6 ++++--
 tools/tool_test.go | 39 +++++++++++++++++++--------------------
 2 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/tools/tool.go b/tools/tool.go
index c76f7f398..b84ad1d81 100644
--- a/tools/tool.go
+++ b/tools/tool.go
@@ -2,9 +2,11 @@ package tools
 
 import (
 	"context"
-	"fmt"
+	"errors"
 )
 
+const ErrInvalidTool = "invalid_tool"
+
 // Tool is a tool for the llm agent to interact with different applications.
 type Tool interface {
 	Name() string
@@ -20,5 +22,5 @@ func (tb *Kit) UseTool(ctx context.Context, toolName string, toolArgs string) (s
 			return tool.Call(ctx, toolArgs)
 		}
 	}
-	return "", fmt.Errorf("invalid tool %v", toolName)
+	return "", errors.New(ErrInvalidTool)
 }
diff --git a/tools/tool_test.go b/tools/tool_test.go
index f5362d8cb..7d300434c 100644
--- a/tools/tool_test.go
+++ b/tools/tool_test.go
@@ -3,45 +3,44 @@ package tools
 import (
 	"context"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
-type SomeTool struct{}
+type someTool struct{}
 
-func (st *SomeTool) Name() string {
+func (st *someTool) Name() string {
 	return "An awesome tool"
 }
 
-func (st *SomeTool) Description() string {
+func (st *someTool) Description() string {
 	return "This tool is awesome"
 }
 
-func (st *SomeTool) Call(ctx context.Context, _ string) (string, error) {
+func (st *someTool) Call(ctx context.Context, _ string) (string, error) {
 	if ctx.Err() != nil {
 		return "", ctx.Err()
 	}
 	return "test", nil
 }
 
-func TestTool(t *testing.T) {
+func TestToolWithTestify(t *testing.T) {
 	t.Parallel()
+	kit := Kit{
+		&someTool{},
+	}
+
+	// Test when the tool exists
 	t.Run("Tool Exists in Kit", func(t *testing.T) {
-		t.Parallel()
-		kit := Kit{
-			&SomeTool{},
-		}
-		_, err := kit.UseTool(context.Background(), "An awesome tool", "test")
-		if err != nil {
-			t.Errorf("Error using tool: %v", err)
-		}
+		result, err := kit.UseTool(context.Background(), "An awesome tool", "test")
+		assert.NoError(t, err)
+		assert.Equal(t, "test", result)
 	})
+
+	// Test when the tool does not exist
 	t.Run("Tool Does Not Exist in Kit", func(t *testing.T) {
-		t.Parallel()
-		kit := Kit{
-			&SomeTool{},
-		}
 		_, err := kit.UseTool(context.Background(), "A tool that does not exist", "test")
-		if err == nil {
-			t.Errorf("Expected error, got nil")
-		}
+		assert.Error(t, err)
+		assert.Equal(t, ErrInvalidTool, err.Error())
 	})
 }