From 34c2df3d79d1534d711d723c5b588cb36b4fb680 Mon Sep 17 00:00:00 2001
From: Pawel <pawelsz125125-github@yahoo.com>
Date: Fri, 23 Sep 2022 00:24:03 -0700
Subject: [PATCH] added context setters

---
 async_transport.go |  4 ++++
 client.go          |  4 ++++
 rollbar.go         |  4 ++++
 rollbar_test.go    | 15 +++++++++++++++
 4 files changed, 27 insertions(+)

diff --git a/async_transport.go b/async_transport.go
index 2dc85e7..b935b43 100644
--- a/async_transport.go
+++ b/async_transport.go
@@ -171,3 +171,7 @@ func (t *AsyncTransport) Close() error {
 func (t *AsyncTransport) setContext(ctx context.Context) {
 	t.ctx = ctx
 }
+
+func (t *AsyncTransport) getContext() context.Context {
+	return t.ctx
+}
diff --git a/client.go b/client.go
index a543412..da6b524 100644
--- a/client.go
+++ b/client.go
@@ -96,6 +96,10 @@ func (c *Client) CaptureTelemetryEvent(eventType, eventlevel string, eventData m
 func (c *Client) SetTelemetry(options ...OptionFunc) {
 	c.Telemetry = NewTelemetry(c.configuration.scrubHeaders, options...)
 }
+func (c *Client) SetContext(ctx context.Context) {
+	c.ctx = ctx
+	c.Transport.setContext(ctx)
+}
 
 // SetEnabled sets whether or not Rollbar is enabled.
 // If this is true then this library works as normal.
diff --git a/rollbar.go b/rollbar.go
index e68270e..7680c30 100644
--- a/rollbar.go
+++ b/rollbar.go
@@ -90,6 +90,10 @@ var DefaultStackTracer StackTracerFunc = func(err error) ([]runtime.Frame, bool)
 	return nil, false
 }
 
+func SetContext(ctx context.Context) {
+	std.SetContext(ctx)
+}
+
 // SetTelemetry sets the telemetry
 func SetTelemetry(options ...OptionFunc) {
 	std.SetTelemetry(options...)
diff --git a/rollbar_test.go b/rollbar_test.go
index 3a2e367..73b3ee3 100644
--- a/rollbar_test.go
+++ b/rollbar_test.go
@@ -10,6 +10,7 @@ import (
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 )
 
 type CustomError struct {
@@ -104,6 +105,20 @@ func TestEverything(t *testing.T) {
 
 type someNonstandardTypeForLogFailing struct{}
 
+func TestSetContext(t *testing.T) {
+	SetToken(os.Getenv("TOKEN"))
+	SetEnvironment("test")
+	ctx, _ := context.WithTimeout(context.Background(), 2*time.Second)
+	SetContext(ctx)
+	if std.ctx != ctx {
+		t.Error("Client ctx must be set")
+	}
+	tr := std.Transport.(*AsyncTransport)
+	if tr.getContext() != ctx {
+		t.Error("Transport ctx must be set")
+	}
+
+}
 func TestEverythingGeneric(t *testing.T) {
 	SetToken(os.Getenv("TOKEN"))
 	SetEnvironment("test")