diff --git a/internals/controllers/network_map/controller/controller.go b/internals/controllers/network_map/controller/controller.go
index d888e45..98f36ee 100644
--- a/internals/controllers/network_map/controller/controller.go
+++ b/internals/controllers/network_map/controller/controller.go
@@ -28,7 +28,7 @@ func NewController(store *db.Store, metrics *appmetrics.AppMetrics, updateChanne
 	}
 }
 
-func (c *Controller) CalculateNetworkMap(accountID string) (*network_map.NetworkMap, error) {
+func (c *Controller) CalculateNetworkMap(ctx context.Context, accountID string) (*network_map.NetworkMap, error) {
 	_, err := c.repo.GetNetworkMapData(accountID)
 	if err != nil {
 		// usually return error
@@ -36,13 +36,13 @@ func (c *Controller) CalculateNetworkMap(accountID string) (*network_map.Network
 
 	// Do calc on data
 
-	log.Tracef("Calculating network map for account on public")
+	log.WithContext(ctx).Tracef("Calculating network map for account on public")
 
 	return &network_map.NetworkMap{}, nil
 }
 
 func (c *Controller) UpdatePeers(ctx context.Context, accountID string) error {
-	_, err := c.CalculateNetworkMap(accountID)
+	_, err := c.CalculateNetworkMap(ctx, accountID)
 	if err != nil {
 		log.Errorf("Failed to calculate network map for account %s: %v", accountID, err)
 		return err
diff --git a/internals/modules/templates/manager/repository_mock.go b/internals/modules/templates/manager/repository_mock.go
index c617857..c36d5d8 100644
--- a/internals/modules/templates/manager/repository_mock.go
+++ b/internals/modules/templates/manager/repository_mock.go
@@ -5,12 +5,13 @@
 package manager
 
 import (
-	context "context"
-	reflect "reflect"
+	"context"
+	"reflect"
 
-	gomock "github.com/golang/mock/gomock"
-	templates "github.com/netbirdio/management-refactor/internals/modules/templates"
-	db "github.com/netbirdio/management-refactor/internals/shared/db"
+	"github.com/golang/mock/gomock"
+
+	"github.com/netbirdio/management-refactor/internals/modules/templates"
+	"github.com/netbirdio/management-refactor/internals/shared/db"
 )
 
 // MockRepository is a mock of Repository interface.