package com.googlecode.objectify.test.util;

import com.google.cloud.datastore.testing.LocalDatastoreHelper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.extension.AfterAllCallback;
import org.junit.jupiter.api.extension.BeforeAllCallback;
import org.junit.jupiter.api.extension.BeforeEachCallback;
import org.junit.jupiter.api.extension.ExtensionContext;
import org.junit.jupiter.api.extension.ExtensionContext.Namespace;

/**
 * Sets up and tears down the Local Datastore emulator, defaults to strong consistency
 */
@RequiredArgsConstructor
@Slf4j
public class LocalDatastoreExtension implements BeforeAllCallback, AfterAllCallback, BeforeEachCallback {

	private static final Namespace NAMESPACE = Namespace.create(LocalDatastoreExtension.class);

	private final double consistency;

	public LocalDatastoreExtension() {
		this(1.0);
	}

	@Override
	public void beforeAll(final ExtensionContext context) throws Exception {
		if (getHelper(context) == null) {
			log.info("Creating new LocalDatastoreHelper");

			final LocalDatastoreHelper helper = LocalDatastoreHelper.create(consistency);
			context.getRoot().getStore(Namespace.GLOBAL).put(LocalDatastoreHelper.class, helper);
			helper.start();
		}

	}

	@Override
	public void afterAll(final ExtensionContext context) throws Exception {
		// We don't really have a way of shutting this down; we need a hook for when the all tests
		// are complete, not just one class. But the java process dies anyways.
//		final LocalDatastoreHelper helper = getHelper(context);
//		helper.stop(Duration.ofSeconds(5));
	}

	@Override
	public void beforeEach(final ExtensionContext context) throws Exception {
		final LocalDatastoreHelper helper = getHelper(context);
		helper.reset();
	}

	/** Get the helper created in beforeAll; it should be global so there will one per test run */
	public static LocalDatastoreHelper getHelper(final ExtensionContext context) {
		return context.getRoot().getStore(Namespace.GLOBAL).get(LocalDatastoreHelper.class, LocalDatastoreHelper.class);
	}
}

@@ -179,65 +179,8 @@
 4j;%0A
-import org.junit.jupiter.api.extension.AfterAllCallback;%0A
 impo
@@ -225,32 +225,32 @@
 oreAllCallback;%0A
+
 import org.junit
@@ -614,26 +614,8 @@
 ack,
- AfterAllCallback,
  Bef
@@ -637,102 +637,8 @@
  %7B%0A%0A
-%09private static final Namespace NAMESPACE = Namespace.create(LocalDatastoreExtension.class);%0A%0A
 %09pri
@@ -1096,363 +1096,8 @@
 %09%7D%0A%0A
-%09@Override%0A%09public void afterAll(final ExtensionContext context) throws Exception %7B%0A%09%09// We don't really have a way of shutting this down; we need a hook for when the all tests%0A%09%09// are complete, not just one class. But the java process dies anyways.%0A//%09%09final LocalDatastoreHelper helper = getHelper(context);%0A//%09%09helper.stop(Duration.ofSeconds(5));%0A%09%7D%0A%0A
 %09@Ov

Clean up LocalDatastoreExtension