From f68e18930052d6880b7f8f87ee8a8b72ad73de54 Mon Sep 17 00:00:00 2001
From: snackmgmg <16898622+snackmgmg@users.noreply.github.com>
Date: Mon, 18 Mar 2024 23:03:43 +0900
Subject: [PATCH] fix: behavior when getting removed callback

---
 callbacks.go            |  7 +++++--
 tests/callbacks_test.go | 22 ++++++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/callbacks.go b/callbacks.go
index 195d1720..aebeed56 100644
--- a/callbacks.go
+++ b/callbacks.go
@@ -154,8 +154,11 @@ func (p *processor) Execute(db *DB) *DB {
 
 func (p *processor) Get(name string) func(*DB) {
 	for i := len(p.callbacks) - 1; i >= 0; i-- {
-		if v := p.callbacks[i]; v.name == name && !v.remove {
-			return v.handler
+		if v := p.callbacks[i]; v.name == name {
+			if !v.remove {
+				return v.handler
+			}
+			return nil
 		}
 	}
 	return nil
diff --git a/tests/callbacks_test.go b/tests/callbacks_test.go
index 4479da4c..59c290c5 100644
--- a/tests/callbacks_test.go
+++ b/tests/callbacks_test.go
@@ -206,3 +206,25 @@ func TestPluginCallbacks(t *testing.T) {
 		t.Errorf("callbacks tests failed, got %v", msg)
 	}
 }
+
+func TestCallbacksRemoveAndGet(t *testing.T) {
+	db, _ := gorm.Open(nil, nil)
+	createCallback := db.Callback().Create()
+
+	err := createCallback.Register("plugin_1_fn1", c1)
+
+	if err != nil {
+		t.Errorf("callbacks tests failed, got error: %v", err)
+	}
+	if cb := createCallback.Get("plugin_1_fn1"); reflect.DeepEqual(cb, c1) {
+		t.Errorf("callbacks tests failed, got: %p, want: %p", cb, c1)
+	}
+
+	err = createCallback.Remove("plugin_1_fn1")
+	if err != nil {
+		t.Errorf("callbacks test failed, got error: %v", err)
+	}
+	if cb := createCallback.Get("plugin_1_fn1"); cb != nil {
+		t.Errorf("callbacks test failed. got: %p, want: nil", cb)
+	}
+}