enhancement for instrumentation of classes under dtest extension

contrib/dtest/src/org/jboss/byteman/contrib/dtest/InstrumentedClass.java

@@ -113,6 +113,29 @@ public void assertMethodCallCount(String message, String methodName, CallCount c
         }
     }
 
+    /**
+     * Checks that the number of known invocations of the given method is specified count.
+     *
+     * @param message the message to print in case of assertion failure.
+     * @param methodName the method name to look for.
+     * @param callCount the expected number of the invocation count.
+     */
+    public void assertMethodCallCount(String message, String methodName, int callCount)
+    {
+    	assertMethodCallCount(message, methodName, new CallCount(callCount, callCount));
+    }
+
+    /**
+     * Checks that the number of known invocations of the given method is specified count.
+     *
+     * @param methodName the method name to look for.
+     * @param callCount the expected number of the invocation count.
+     */
+    public void assertMethodCallCount(String methodName, int callCount)
+    {
+    	assertMethodCallCount(null, methodName, new CallCount(callCount, callCount));
+    }
+
     /**
      * Checks that the given method has been called at least once on each known instance of the class.
      * Uses junit internally, hence expect the normal exception throwing in case of failure.

contrib/dtest/src/org/jboss/byteman/contrib/dtest/InstrumentedInstance.java

@@ -89,7 +89,13 @@ public int getInvocationCount(String methodName)
     public void assertMethodCallCount(String message, String methodName, CallCount callCount)
     {
         int invocationCount = getInvocationCount(methodName);
-        assertTrue((message == null ? "" : message)+" - required minimum call count "+callCount.getMin()+" but was "+invocationCount, callCount.getMin() <= invocationCount);
-        assertTrue((message == null ? "" : message)+" - required maximum call count "+callCount.getMax()+" but was "+invocationCount, callCount.getMax() >= invocationCount);
+
+        String assertInfo = (message == null ? "" : message + " - ") + String.format("Method %s#%s ", className, methodName);
+        if(callCount.getMin() == callCount.getMax()) {
+        	assertTrue(assertInfo + "required call count " + callCount.getMin() + " but was " + invocationCount, callCount.getMin() == invocationCount);
+        } else {
+        	assertTrue(assertInfo + "required minimum call count " + callCount.getMin() + " but was " + invocationCount, callCount.getMin() <= invocationCount);
+        	assertTrue(assertInfo + "required maximum call count " + callCount.getMax() + " but was " + invocationCount, callCount.getMax() >= invocationCount);
+        }
     }
 }

contrib/dtest/src/org/jboss/byteman/contrib/dtest/Instrumentor.java

@@ -111,29 +111,55 @@ public InstrumentedClass instrumentClass(Class clazz) throws Exception
     public InstrumentedClass instrumentClass(Class clazz, Set<String> methodNames) throws Exception
     {
         String className = clazz.getCanonicalName();
+
+        Set<String> methodNamesToInstrument = new HashSet<String>();
+
+        for(Method method : clazz.getDeclaredMethods()) {
+        	String declaredMethodName = method.getName();
+        	if(methodNames == null || methodNames.contains(declaredMethodName)) {
+        		methodNamesToInstrument.add(declaredMethodName);
+        	}
+        }
+
+        return instrumentClass(className, methodNamesToInstrument);
+    }
+
+    /**
+     * Add method tracing rules to the specified class name.<br>
+     * If a null set of method names is supplied, {@link NullPointerException} is thrown.
+     *
+     * @param className the class name to instrument.
+     * @param methodNames the selection of methods to instrument.
+     * @return a local proxy for the instrumentation.
+     * @throws NullPointerException in case of methodNames parameter is null
+     * @throws Exception in case of failure.
+     */
+    public InstrumentedClass instrumentClass(String className, Set<String> methodNames) throws Exception
+    {
+    	if(methodNames == null) {
+    		throw new NullPointerException("methodNames");
+    	}
+
         Set<String> instrumentedMethods = new HashSet<String>();
 
         StringBuilder ruleScriptBuilder = new StringBuilder();
-        for(Method method : clazz.getDeclaredMethods()) {
-            String methodName = method.getName();
+        for(String methodName : methodNames) {
 
             if(instrumentedMethods.contains(methodName)) {
               // do not add two identical rules for methods which differ by parameters
               continue;
             }
 
-            if(methodNames == null || methodNames.contains(methodName)) {
 
-                String ruleName = this.getClass().getCanonicalName()+"_"+className+"_"+methodName+"_remotetrace_entry";
+            String ruleName = this.getClass().getCanonicalName()+"_"+className+"_"+methodName+"_remotetrace_entry";
 
-                RuleBuilder ruleBuilder = new RuleBuilder(ruleName);
-                ruleBuilder.onClass(className).inMethod(methodName).atEntry();
-                ruleBuilder.usingHelper(BytemanTestHelper.class);
-                ruleBuilder.doAction("setTriggering(false), debug(\"firing "+ruleName+"\", $0), remoteTrace(\""+className+"\", \""+methodName+"\", $*)");
-                ruleScriptBuilder.append(ruleBuilder.toString());
-
-                instrumentedMethods.add(methodName);
-            }
+            RuleBuilder ruleBuilder = new RuleBuilder(ruleName);
+            ruleBuilder.onClass(className).inMethod(methodName).atEntry();
+            ruleBuilder.usingHelper(BytemanTestHelper.class);
+            ruleBuilder.doAction("setTriggering(false), debug(\"firing "+ruleName+"\", $0), remoteTrace(\""+className+"\", \""+methodName+"\", $*)");
+            ruleScriptBuilder.append(ruleBuilder.toString());
+
+            instrumentedMethods.add(methodName);
         }
 
         String scriptString = ruleScriptBuilder.toString();