From 0d6a8cf73b3f5b172315b9eb7e5e18824e32b85c Mon Sep 17 00:00:00 2001
From: Scott Opell <me@scottopell.com>
Date: Fri, 31 May 2024 11:26:25 -0400
Subject: [PATCH] fixes tests?

---
 src/main/java/org/datadog/jmxfetch/App.java   | 16 ++++++++------
 .../java/org/datadog/jmxfetch/Instance.java   | 22 ++++++++++++-------
 .../java/org/datadog/jmxfetch/TestApp.java    |  5 -----
 3 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/src/main/java/org/datadog/jmxfetch/App.java b/src/main/java/org/datadog/jmxfetch/App.java
index 77ad357e9..9e40f1ec7 100644
--- a/src/main/java/org/datadog/jmxfetch/App.java
+++ b/src/main/java/org/datadog/jmxfetch/App.java
@@ -133,10 +133,12 @@ public App(final AppConfig appConfig) {
         }
         this.configs = getConfigs(this.appConfig);
 
-//        if (this.appConfig.getJmxfetchTelemetry()) {
+        this.appTelemetry = new AppTelemetry();
+
+        if (this.appConfig.getJmxfetchTelemetry()) {
             log.info("Enabling JMX Fetch Telemetry");
-            this.initTelemetryBean();
-//        }
+            this.registerTelemetryBean(this.appTelemetry);
+        }
     }
 
     private ObjectName getAppTelemetryBeanName() {
@@ -156,9 +158,8 @@ private ObjectName getAppTelemetryBeanName() {
         return appTelemetryBeanName;
     }
 
-    private void initTelemetryBean() {
+    private void registerTelemetryBean(AppTelemetry bean) {
         MBeanServer mbs =  ManagementFactory.getPlatformMBeanServer();
-        AppTelemetry bean = new AppTelemetry();
         ObjectName appTelemetryBeanName = getAppTelemetryBeanName();
         if (appTelemetryBeanName == null) {
             return;
@@ -174,7 +175,6 @@ private void initTelemetryBean() {
                 appTelemetryBeanName.toString(), e);
         }
 
-        this.appTelemetry = bean;
         return;
     }
 
@@ -523,7 +523,9 @@ void start() {
     }
 
     void stop() {
-        this.teardownTelemetry();
+        if (this.appConfig.getJmxfetchTelemetry()) {
+            this.teardownTelemetry();
+        }
         this.collectionProcessor.stop();
         this.recoveryProcessor.stop();
     }
diff --git a/src/main/java/org/datadog/jmxfetch/Instance.java b/src/main/java/org/datadog/jmxfetch/Instance.java
index 9da8de750..12f95ad83 100644
--- a/src/main/java/org/datadog/jmxfetch/Instance.java
+++ b/src/main/java/org/datadog/jmxfetch/Instance.java
@@ -240,7 +240,10 @@ public Instance(
             log.info("collect_default_jvm_metrics is false - not collecting default JVM metrics");
         }
 
-        instanceTelemetryBean = createInstanceTelemetryBean();
+        instanceTelemetryBean = new InstanceTelemetry();
+        if (appConfig.getJmxfetchTelemetry()) {
+            registerTelemetryBean(instanceTelemetryBean);
+        }
     }
 
     private ObjectName getObjName(String domain,String instance)
@@ -248,10 +251,9 @@ private ObjectName getObjName(String domain,String instance)
         return new ObjectName(domain + ":target_instance=" + ObjectName.quote(instance));
     }
 
-    private InstanceTelemetry createInstanceTelemetryBean() {
+    private InstanceTelemetry registerTelemetryBean(InstanceTelemetry bean) {
+        log.debug("Registering jmx bean for instance: " + this.getCheckName());
         mbs =  ManagementFactory.getPlatformMBeanServer();
-        InstanceTelemetry bean = new InstanceTelemetry();
-        log.debug("Created jmx bean for instance: " + this.getCheckName());
 
         try {
             instanceTelemetryBeanName = getObjName(appConfig.getJmxfetchTelemetryDomain(),
@@ -516,8 +518,8 @@ public List<Metric> getMetrics() throws IOException {
                     + " With beans fetched = " + instanceTelemetryBean.getBeansFetched()
                     + " top attributes = " + instanceTelemetryBean.getTopLevelAttributeCount()
                     + " metrics = " + instanceTelemetryBean.getMetricCount()
-                    + " wildcard domain query count = " 
-                    + instanceTelemetryBean.getWildcardDomainQueryCount() 
+                    + " wildcard domain query count = "
+                    + instanceTelemetryBean.getWildcardDomainQueryCount()
                     + " bean match ratio = " + instanceTelemetryBean.getBeanMatchRatio());
         }
         return metrics;
@@ -710,7 +712,7 @@ private void getMatchingAttributes() throws IOException {
                     reporter.displayNonMatchingAttributeName(jmxAttribute);
                 }
                 if (jmxAttribute.getMatchingConf() != null) {
-                    attributeMatched = true; 
+                    attributeMatched = true;
                 }
             }
             if (attributeMatched) {
@@ -718,7 +720,7 @@ private void getMatchingAttributes() throws IOException {
             }
         }
         if (instanceTelemetryBean != null) {
-            instanceTelemetryBean.setBeanMatchRatio((double) 
+            instanceTelemetryBean.setBeanMatchRatio((double)
                                   beansWithAttributeMatch / beans.size());
         }
         log.info("Found {} matching attributes", matchingAttributes.size());
@@ -837,6 +839,10 @@ public boolean isLimitReached() {
     }
 
     private void cleanupTelemetryBean() {
+        if (!appConfig.getJmxfetchTelemetry()) {
+            // If telemetry is not enabled, no need to unregister the bean
+            return;
+        }
         try {
             mbs.unregisterMBean(instanceTelemetryBeanName);
             log.debug("Successfully unregistered bean for instance: " + this.getCheckName());
diff --git a/src/test/java/org/datadog/jmxfetch/TestApp.java b/src/test/java/org/datadog/jmxfetch/TestApp.java
index 23f180cf6..028f5c75b 100644
--- a/src/test/java/org/datadog/jmxfetch/TestApp.java
+++ b/src/test/java/org/datadog/jmxfetch/TestApp.java
@@ -25,8 +25,6 @@ public class TestApp extends TestCommon {
     @Test
     public void testBeanRegexTags() throws Exception {
         log.info("********************\nSTARTING\n********************");
-        // When we enable JMXFetch telemetry
-        when(appConfig.getJmxfetchTelemetry()).thenReturn(true);
 
 
         log.info("********************\nregisterMBean\n********************");
@@ -64,9 +62,6 @@ public void testBeanRegexTags() throws Exception {
     /** Tag metrics with MBeans parameters. */
     @Test
     public void testBeanTags() throws Exception {
-        // When we enable JMXFetch telemetry
-        when(appConfig.getJmxfetchTelemetry()).thenReturn(true);
-
         // We expose a few metrics through JMX
         registerMBean(
                 new SimpleTestJavaApp(),