diff --git a/agent/main.go b/agent/main.go index 7974669f8..1db69f27d 100644 --- a/agent/main.go +++ b/agent/main.go @@ -216,7 +216,7 @@ func enableExpFeatureIfNeeded() { oas.GetOasGeneratorInstance().Start() } if config.Config.ServiceMap { - servicemap.GetInstance().SetConfig(config.Config) + servicemap.GetInstance().Enable() } elastic.GetInstance().Configure(config.Config.Elastic) } diff --git a/agent/pkg/controllers/service_map_controller_test.go b/agent/pkg/controllers/service_map_controller_test.go index c6b372406..9602efd7e 100644 --- a/agent/pkg/controllers/service_map_controller_test.go +++ b/agent/pkg/controllers/service_map_controller_test.go @@ -11,7 +11,6 @@ import ( "github.com/gin-gonic/gin" "github.com/stretchr/testify/suite" - "github.com/up9inc/mizu/shared" tapApi "github.com/up9inc/mizu/tap/api" ) @@ -59,9 +58,7 @@ type ServiceMapControllerSuite struct { func (s *ServiceMapControllerSuite) SetupTest() { s.c = NewServiceMapController() - s.c.service.SetConfig(&shared.MizuAgentConfig{ - ServiceMap: true, - }) + s.c.service.Enable() s.c.service.NewTCPEntry(TCPEntryA, TCPEntryB, ProtocolHttp) s.w = httptest.NewRecorder() diff --git a/agent/pkg/servicemap/servicemap.go b/agent/pkg/servicemap/servicemap.go index 3537e9855..3f7359feb 100644 --- a/agent/pkg/servicemap/servicemap.go +++ b/agent/pkg/servicemap/servicemap.go @@ -3,7 +3,6 @@ package servicemap import ( "sync" - "github.com/up9inc/mizu/shared" "github.com/up9inc/mizu/shared/logger" tapApi "github.com/up9inc/mizu/tap/api" ) @@ -26,13 +25,13 @@ func GetInstance() ServiceMap { } type serviceMap struct { - config *shared.MizuAgentConfig + enabled bool graph *graph entriesProcessed int } type ServiceMap interface { - SetConfig(config *shared.MizuAgentConfig) + Enable() IsEnabled() bool NewTCPEntry(source *tapApi.TCP, destination *tapApi.TCP, protocol *tapApi.Protocol) GetStatus() ServiceMapStatus @@ -46,7 +45,7 @@ type ServiceMap interface { func newServiceMap() *serviceMap { return &serviceMap{ - config: nil, + enabled: false, entriesProcessed: 0, graph: newDirectedGraph(), } @@ -156,15 +155,12 @@ func (s *serviceMap) addEdge(u, v *entryData, p *tapApi.Protocol) { s.entriesProcessed++ } -func (s *serviceMap) SetConfig(config *shared.MizuAgentConfig) { - s.config = config +func (s *serviceMap) Enable() { + s.enabled = true } func (s *serviceMap) IsEnabled() bool { - if s.config != nil && s.config.ServiceMap { - return true - } - return false + return s.enabled } func (s *serviceMap) NewTCPEntry(src *tapApi.TCP, dst *tapApi.TCP, p *tapApi.Protocol) { diff --git a/agent/pkg/servicemap/servicemap_test.go b/agent/pkg/servicemap/servicemap_test.go index 9a0de04b6..6dc2126e4 100644 --- a/agent/pkg/servicemap/servicemap_test.go +++ b/agent/pkg/servicemap/servicemap_test.go @@ -6,7 +6,6 @@ import ( "testing" "github.com/stretchr/testify/suite" - "github.com/up9inc/mizu/shared" tapApi "github.com/up9inc/mizu/tap/api" ) @@ -96,9 +95,7 @@ func (s *ServiceMapDisabledSuite) SetupTest() { func (s *ServiceMapEnabledSuite) SetupTest() { s.instance = GetInstance() - s.instance.SetConfig(&shared.MizuAgentConfig{ - ServiceMap: true, - }) + s.instance.Enable() } func (s *ServiceMapDisabledSuite) TestServiceMapInstance() {