1 /*
2  * Copyright (C) 2018 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License"); you
5  * may not use this file except in compliance with the License. You may
6  * obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
13  * implied. See the License for the specific language governing
14  * permissions and limitations under the License.
15  */
16 
17 package com.android.vts.util;
18 
19 import com.google.cloud.datastore.Datastore;
20 import com.googlecode.objectify.ObjectifyFactory;
21 import com.googlecode.objectify.ObjectifyService;
22 import com.googlecode.objectify.util.Closeable;
23 import org.junit.jupiter.api.extension.AfterEachCallback;
24 import org.junit.jupiter.api.extension.BeforeEachCallback;
25 import org.junit.jupiter.api.extension.ExtensionContext;
26 import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
27 
28 /** Sets up and tears down the GAE local unit test harness environment */
29 public class ObjectifyExtension implements BeforeEachCallback, AfterEachCallback {
30 
31     private static final Namespace NAMESPACE = Namespace.create(ObjectifyExtension.class);
32 
33     @Override
beforeEach(final ExtensionContext context)34     public void beforeEach(final ExtensionContext context) throws Exception {
35         final Datastore datastore =
36                 LocalDatastoreExtension.getHelper(context).getOptions().getService();
37 
38         ObjectifyService.init(new ObjectifyFactory(datastore));
39 
40         final Closeable rootService = ObjectifyService.begin();
41 
42         context.getStore(NAMESPACE).put(Closeable.class, rootService);
43     }
44 
45     @Override
afterEach(final ExtensionContext context)46     public void afterEach(final ExtensionContext context) throws Exception {
47         final Closeable rootService =
48                 context.getStore(NAMESPACE).get(Closeable.class, Closeable.class);
49 
50         rootService.close();
51     }
52 }
53