mockit.internal.state.MockFixture 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-2014 Rogério Liesenfeld
 * This file is subject to the terms of the MIT license (see LICENSE.txt).
 */
package mockit.internal.state;

import java.lang.reflect.*;
import java.util.*;
import java.util.Map.*;

import org.jetbrains.annotations.*;

import mockit.internal.*;
import mockit.internal.expectations.mocking.*;
import mockit.internal.util.*;

/**
 * Holds data about redefined real classes and their corresponding mock classes (if any), and provides methods to
 * add/remove such state both from this instance and from other state holders with associated data.
 */
public final class MockFixture
{
   /**
    * Similar to {@code redefinedClasses}, but for classes modified by a {@code ClassFileTransformer} such as the
    * {@code CaptureTransformer}, and containing the pre-transform bytecode instead of the modified one.
    */
   private final Map transformedClasses = new HashMap(2);

   /**
    * Real classes currently redefined in the running JVM and their current (modified) bytecodes.
    * 
    * The keys in the map allow each redefined real class to be later restored to a previous definition.
    * 

    * The modified bytecode arrays in the map allow a new redefinition to be made on top of the current redefinition
    * (in the case of the Mockups API), or to restore the class to a previous definition (provided the map is copied
    * between redefinitions of the same class).
    */
   private final Map, byte[]> redefinedClasses = new IdentityHashMap, byte[]>(8);

   /**
    * Subset of all currently redefined classes which contain one or more native methods.
    * 

    * This is needed because in order to restore such methods it is necessary (for some classes) to re-register them
    * with the JVM.
    *
    * @see #reregisterNativeMethodsForRestoredClass(Class)
    */
   private final Set redefinedClassesWithNativeMethods = new HashSet();

   /**
    * Maps redefined real classes to the internal name of the corresponding mock classes, when it's the case.
    * 
    * This allows any global state associated to a mock class to be discarded when the corresponding real class is
    * later restored to its original definition.
    */
   private final Map, String> realClassesToMockClasses = new IdentityHashMap, String>(8);

   private final List> mockedClasses = new ArrayList>();
   private final Map mockedTypesAndInstances = new IdentityHashMap();

   // Methods to add/remove transformed/redefined classes /////////////////////////////////////////////////////////////

   public void addTransformedClass(@NotNull String className, @NotNull byte[] pretransformClassfile)
   {
      transformedClasses.put(className, pretransformClassfile);
   }

   public void addRedefinedClass(
      @Nullable String mockClassInternalName, @NotNull Class redefinedClass, @NotNull byte[] modifiedClassfile)
   {
      if (mockClassInternalName != null) {
         String previousNames = realClassesToMockClasses.put(redefinedClass, mockClassInternalName);

         if (previousNames != null) {
            realClassesToMockClasses.put(redefinedClass, previousNames + ' ' + mockClassInternalName);
         }
      }

      addRedefinedClass(redefinedClass, modifiedClassfile);
   }

   public void addRedefinedClass(@NotNull Class redefinedClass, @NotNull byte[] modifiedClassfile)
   {
      redefinedClasses.put(redefinedClass, modifiedClassfile);
   }

   public void registerMockedClass(@NotNull Class mockedType)
   {
      if (!mockedClasses.contains(mockedType) && !GeneratedClasses.isGeneratedImplementationClass(mockedType)) {
         mockedClasses.add(mockedType);
      }
   }

   public boolean isInstanceOfMockedClass(@NotNull Object mockedInstance)
   {
      Class mockedClass = mockedInstance.getClass();
      int n = mockedClasses.size();

      for (int i = 0; i < n; i++) {
         Class mockedType = mockedClasses.get(i);

         if (mockedType == mockedClass || mockedType.isAssignableFrom(mockedClass)) {
            return true;
         }
      }

      return false;
   }

   public void registerInstanceFactoryForMockedType(
      @NotNull Class mockedType, @NotNull InstanceFactory mockedInstanceFactory)
   {
      registerMockedClass(mockedType);
      mockedTypesAndInstances.put(mockedType, mockedInstanceFactory);
   }

   @Nullable public InstanceFactory findInstanceFactory(@NotNull Type mockedType)
   {
      Class mockedClass = Utilities.getClassType(mockedType);

      if (mockedClass.isInterface() || Modifier.isAbstract(mockedClass.getModifiers())) {
         for (Entry entry : mockedTypesAndInstances.entrySet()) {
            Type registeredMockedType = entry.getKey();
            Class registeredMockedClass = Utilities.getClassType(registeredMockedType);
            Class baseType = GeneratedClasses.getMockedClassOrInterfaceType(registeredMockedClass);

            if (baseType == mockedClass) {
               return entry.getValue();
            }
         }

         return null;
      }

      return mockedTypesAndInstances.get(mockedType);
   }

   public void restoreAndRemoveRedefinedClasses(@Nullable Set> desiredClasses)
   {
      Set> classesToRestore = desiredClasses == null ? redefinedClasses.keySet() : desiredClasses;
      RedefinitionEngine redefinitionEngine = new RedefinitionEngine();

      for (Class redefinedClass : classesToRestore) {
         redefinitionEngine.restoreOriginalDefinition(redefinedClass);
         restoreDefinition(redefinedClass);
         discardStateForCorrespondingMockClassIfAny(redefinedClass);
      }

      if (desiredClasses == null) {
         redefinedClasses.clear();
      }
      else {
         redefinedClasses.keySet().removeAll(desiredClasses);
      }
   }

   private void restoreDefinition(@NotNull Class redefinedClass)
   {
      if (redefinedClassesWithNativeMethods.contains(redefinedClass.getName())) {
         reregisterNativeMethodsForRestoredClass(redefinedClass);
      }

      mockedTypesAndInstances.remove(redefinedClass);
      mockedClasses.remove(redefinedClass);
   }

   private void discardStateForCorrespondingMockClassIfAny(@NotNull Class redefinedClass)
   {
      String mockClassesInternalNames = realClassesToMockClasses.remove(redefinedClass);
      TestRun.getMockStates().removeClassState(redefinedClass, mockClassesInternalNames);
   }

   void restoreTransformedClasses(@NotNull Set previousTransformedClasses)
   {
      if (!transformedClasses.isEmpty()) {
         Set classesToRestore;

         if (previousTransformedClasses.isEmpty()) {
            classesToRestore = transformedClasses.keySet();
         }
         else {
            classesToRestore = getTransformedClasses();
            classesToRestore.removeAll(previousTransformedClasses);
         }

         if (!classesToRestore.isEmpty()) {
            restoreAndRemoveTransformedClasses(classesToRestore);
         }
      }
   }

   private void restoreAndRemoveTransformedClasses(@NotNull Set classesToRestore)
   {
      RedefinitionEngine redefinitionEngine = new RedefinitionEngine();

      for (String transformedClassName : classesToRestore) {
         byte[] definitionToRestore = transformedClasses.get(transformedClassName);
         redefinitionEngine.restoreToDefinition(transformedClassName, definitionToRestore);
      }

      transformedClasses.keySet().removeAll(classesToRestore);
   }

   void restoreRedefinedClasses(@NotNull Map previousDefinitions)
   {
      RedefinitionEngine redefinitionEngine = new RedefinitionEngine();
      Iterator, byte[]>> itr = redefinedClasses.entrySet().iterator();

      while (itr.hasNext()) {
         Entry, byte[]> entry = itr.next();
         Class redefinedClass = entry.getKey();
         byte[] currentDefinition = entry.getValue();
         byte[] previousDefinition = previousDefinitions.get(redefinedClass);

         if (currentDefinition != previousDefinition) {
            redefinitionEngine.restoreDefinition(redefinedClass, previousDefinition);

            if (previousDefinition == null) {
               restoreDefinition(redefinedClass);
               discardStateForCorrespondingMockClassIfAny(redefinedClass);
               itr.remove();
            }
            else {
               entry.setValue(previousDefinition);
            }
         }
      }
   }

   // Methods that deal with redefined native methods /////////////////////////////////////////////////////////////////

   public void addRedefinedClassWithNativeMethods(@NotNull String redefinedClassInternalName)
   {
      redefinedClassesWithNativeMethods.add(redefinedClassInternalName.replace('/', '.'));
   }

   private void reregisterNativeMethodsForRestoredClass(@NotNull Class realClass)
   {
      Method registerNatives = null;

      try {
         registerNatives = realClass.getDeclaredMethod("registerNatives");
      }
      catch (NoSuchMethodException ignore) {
         try {
            registerNatives = realClass.getDeclaredMethod("initIDs");
         }
         catch (NoSuchMethodException alsoIgnore) {
            // OK
         }
      }

      if (registerNatives != null) {
         try {
            registerNatives.setAccessible(true);
            registerNatives.invoke(null);
         }
         catch (IllegalAccessException ignore) {
            // Won't happen.
         }
         catch (InvocationTargetException ignore) {
            // Shouldn't happen either.
         }
      }

      // OK, although another solution will be required for this particular class if it requires
      // natives to be explicitly registered again (not all do, such as java.lang.Float).
   }

   // Getter methods for the maps of transformed/redefined classes ////////////////////////////////////////////////////

   @NotNull public Set getTransformedClasses() { return new HashSet(transformedClasses.keySet()); }

   @NotNull public Map, byte[]> getRedefinedClasses()
   {
      return new HashMap, byte[]>(redefinedClasses);
   }

   @Nullable public byte[] getRedefinedClassfile(@NotNull Class redefinedClass)
   {
      return redefinedClasses.get(redefinedClass);
   }

   public boolean containsRedefinedClass(@NotNull Class redefinedClass)
   {
      return redefinedClasses.containsKey(redefinedClass);
   }
}