package com.braintreepayments.testutils;

import android.app.Activity;
import android.app.KeyguardManager;
import android.content.Context;
import android.support.test.espresso.intent.Intents;
import android.support.test.rule.ActivityTestRule;

import static android.support.test.InstrumentationRegistry.getTargetContext;
import static com.braintreepayments.testutils.SharedPreferencesHelper.getSharedPreferences;

public class BraintreeActivityTestRule<T extends Activity> extends ActivityTestRule<T> {

    public BraintreeActivityTestRule(Class<T> activityClass) {
        super(activityClass);
    }

    public BraintreeActivityTestRule(Class<T> activityClass, boolean initialTouchMode,
            boolean launchActivity) {
        super(activityClass, initialTouchMode, launchActivity);
    }

    @Override
    protected void beforeActivityLaunched() {
        super.beforeActivityLaunched();
        getSharedPreferences().edit().clear().commit();

        ((KeyguardManager) getTargetContext().getSystemService(Context.KEYGUARD_SERVICE))
                .newKeyguardLock("BraintreeActivityTestRule")
                .disableKeyguard();
    }

    @Override
    protected void afterActivityLaunched() {
        Intents.init();
        super.afterActivityLaunched();
    }

    @Override
    protected void afterActivityFinished() {
        super.afterActivityFinished();

        Intents.release();

        getSharedPreferences().edit().clear().commit();
    }
}

@@ -100,16 +100,65 @@
 anager;%0A
+import android.app.KeyguardManager.KeyguardLock;%0A
 import a
@@ -546,16 +546,57 @@
 e%3CT%3E %7B%0A%0A
+    private KeyguardLock mKeyguardLock;%0A%0A
     publ
@@ -680,16 +680,32 @@
 Class);%0A
+        init();%0A
     %7D%0A%0A 
@@ -900,110 +900,52 @@
     
-%7D%0A%0A
     
-@Override%0A    protected void beforeActivityLaunched() %7B
+init();
 %0A    
+%7D%0A%0A
     
-super.beforeActivityLaunched();
+private void init() %7B
 %0A   
@@ -1005,16 +1005,32 @@
 %0A       
+ mKeyguardLock =
  ((Keygu
@@ -1165,16 +1165,17 @@
 stRule%22)
+;
 %0A       
@@ -1175,24 +1175,29 @@
         
-        
+mKeyguardLock
 .disable
@@ -1439,24 +1439,42 @@
 inished();%0A%0A
+        try %7B%0A    
         Inte
@@ -1487,16 +1487,67 @@
 lease();
+%0A        %7D catch (IllegalStateException ignored) %7B%7D
 %0A%0A      
@@ -1591,21 +1591,64 @@
 ().commit();
+%0A%0A        mKeyguardLock.reenableKeyguard();
 %0A    %7D%0A%7D%0A

Fix incorrect test failures and Keyguard leaks