mockit.internal.capturing.CaptureTransformer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jmockit Show documentation
JMockit is a Java toolkit for developer (unit/integration) testing. It contains mocking APIs and other tools, supporting both JUnit and TestNG. The mocking APIs allow all kinds of Java code, without testability restrictions, to be tested in isolation from selected dependencies.
The newest version!
/*
 * Copyright (c) 2006-2013 Rogério Liesenfeld
 * This file is subject to the terms of the MIT license (see LICENSE.txt).
 */
package mockit.internal.capturing;

import java.lang.instrument.*;
import java.security.*;
import java.util.*;

import org.jetbrains.annotations.*;

import mockit.external.asm4.*;
import mockit.internal.*;
import mockit.internal.state.*;
import mockit.internal.util.*;

import static mockit.internal.util.GeneratedClasses.*;

final class CaptureTransformer implements ClassFileTransformer
{
   @NotNull private final String capturedType;
   @NotNull private final CaptureOfImplementations modifierFactory;
   @Nullable private final Map transformedClasses;
   private boolean inactive;

   CaptureTransformer(
      @NotNull CapturedType metadata, @NotNull CaptureOfImplementations modifierFactory,
      boolean registerTransformedClasses)
   {
      capturedType = Type.getInternalName(metadata.baseType);
      this.modifierFactory = modifierFactory;
      transformedClasses = registerTransformedClasses ? new HashMap(2) : null;
   }

   void deactivate()
   {
      inactive = true;

      if (transformedClasses != null) {
         RedefinitionEngine redefinitionEngine = new RedefinitionEngine();

         for (Map.Entry classNameAndOriginalBytecode : transformedClasses.entrySet()) {
            String className = classNameAndOriginalBytecode.getKey();
            byte[] originalBytecode = classNameAndOriginalBytecode.getValue();
            redefinitionEngine.restoreToDefinition(className, originalBytecode);
         }

         transformedClasses.clear();
      }
   }

   @Nullable
   public byte[] transform(
      @Nullable ClassLoader loader, @NotNull String internalClassName, @Nullable Class classBeingRedefined,
      @Nullable ProtectionDomain protectionDomain, @NotNull byte[] classfileBuffer)
   {
      if (
         inactive || classBeingRedefined != null || TestRun.getCurrentTestInstance() == null ||
         internalClassName.startsWith("mockit/internal/"))
      {
         return null;
      }

      ClassReader cr = new ClassReader(classfileBuffer);
      SuperTypeCollector superTypeCollector = new SuperTypeCollector(loader);

      try {
         cr.accept(superTypeCollector, ClassReader.SKIP_DEBUG);
      }
      catch (VisitInterruptedException ignore) {
         if (superTypeCollector.classExtendsCapturedType && !isGeneratedClass(internalClassName)) {
            String className = internalClassName.replace('/', '.');
            return modifyAndRegisterClass(loader, className, cr);
         }
      }

      return null;
   }

   @NotNull
   private byte[] modifyAndRegisterClass(
      @Nullable ClassLoader loader, @NotNull String className, @NotNull ClassReader cr)
   {
      ClassVisitor modifier = modifierFactory.createModifier(loader, cr, capturedType);
      cr.accept(modifier, 0);

      byte[] originalBytecode = cr.b;

      if (transformedClasses == null) {
         TestRun.mockFixture().addTransformedClass(className, originalBytecode);
      }
      else {
         transformedClasses.put(className, originalBytecode);
      }

      return modifier.toByteArray();
   }

   private final class SuperTypeCollector extends ClassVisitor
   {
      @Nullable private final ClassLoader loader;
      boolean classExtendsCapturedType;

      private SuperTypeCollector(@Nullable ClassLoader loader) { this.loader = loader; }

      @Override
      public void visit(
         int version, int access, @NotNull String name, @Nullable String signature, @Nullable String superName,
         @Nullable String[] interfaces)
      {
         classExtendsCapturedType = false;

         if (capturedType.equals(superName)) {
            classExtendsCapturedType = true;
         }
         else if (interfaces != null && interfaces.length > 0) {
            for (String itfc : interfaces) {
               if (capturedType.equals(itfc)) {
                  classExtendsCapturedType = true;
                  break;
               }
            }
         }

         if (superName != null && !classExtendsCapturedType && !"java/lang/Object".equals(superName)) {
            ClassReader cr = ClassFile.createClassFileReader(loader, superName);
            cr.accept(this, ClassReader.SKIP_DEBUG);
         }

         throw VisitInterruptedException.INSTANCE;
      }
   }
}