From c5fbda3f6faeb8c5a21b59a7da5113e80020929e Mon Sep 17 00:00:00 2001
From: Ali <bincyber@users.noreply.github.com>
Date: Sat, 13 Mar 2021 17:34:07 -0500
Subject: [PATCH] Added tests

---
 context_test.go | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)
 create mode 100644 context_test.go

diff --git a/context_test.go b/context_test.go
new file mode 100644
index 00000000..c4885925
--- /dev/null
+++ b/context_test.go
@@ -0,0 +1,34 @@
+package gorm
+
+import (
+	"context"
+	"reflect"
+	"testing"
+)
+
+func TestContext_success(t *testing.T) {
+	db := &DB{}
+
+	ctx := WithContext(context.Background(), db)
+	extractedDB, err := FromContext(ctx)
+
+	if err != nil {
+		t.Errorf("expected err to be nil. Got: %v", err)
+	}
+
+	if !reflect.DeepEqual(db, extractedDB) {
+		t.Errorf("db and extractedDB are not the same")
+	}
+}
+
+func TestContext_failure(t *testing.T) {
+	extractedDB, err := FromContext(context.Background())
+
+	if extractedDB != nil {
+		t.Errorf("expected extractedDB to nil. Got: %v", extractedDB)
+	}
+
+	if err != ErrDBNotFoundInContext {
+		t.Errorf("expected err to be ErrDBNotFoundInContext. Got: %v", err)
+	}
+}