001/**
002 * Copyright (C) 2006-2022 Talend Inc. - www.talend.com
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 * http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package org.talend.sdk.component.junit.environment.builtin.beam;
017
018import static org.jboss.shrinkwrap.resolver.api.maven.ScopeType.RUNTIME;
019import static org.jboss.shrinkwrap.resolver.api.maven.coordinate.MavenDependencies.createDependency;
020
021import java.lang.annotation.Annotation;
022import java.lang.reflect.Method;
023import java.util.function.Consumer;
024
025import org.jboss.shrinkwrap.resolver.api.maven.coordinate.MavenDependency;
026import org.talend.sdk.component.junit.environment.ClassLoaderEnvironment;
027
028import lombok.extern.slf4j.Slf4j;
029
030@Slf4j
031public abstract class BeamEnvironment extends ClassLoaderEnvironment {
032
033    private String beamVersion;
034
035    private String kitVersion;
036
037    @Override
038    protected AutoCloseable doStart(final Class<?> clazz, final Annotation[] annotations) {
039        beamVersion = System.getProperty("talend.junit.beam.version", Versions.BEAM_VERSION);
040        kitVersion = System.getProperty("talend.junit.kit.version", Versions.KIT_VERSION);
041        resetBeamCache(false);
042        final AutoCloseable delegate = super.doStart(clazz, annotations);
043        return () -> {
044            resetBeamCache(true);
045            delegate.close();
046        };
047    }
048
049    private void resetBeamCache(final boolean highLevelLog) {
050        try { // until resetCache() is part of beam do it the hard way
051            final ClassLoader loader = Thread.currentThread().getContextClassLoader();
052            final Class<?> pof = loader.loadClass("org.apache.beam.sdk.options.PipelineOptionsFactory");
053
054            final Method initializeRegistry = pof.getDeclaredMethod("resetCache");
055            initializeRegistry.setAccessible(true);
056            initializeRegistry.invoke(null);
057        } catch (final NoClassDefFoundError | Exception ex) {
058            final Consumer<String> logger = highLevelLog ? log::warn : log::debug;
059            logger.accept(ex.getMessage());
060        }
061    }
062
063    @Override
064    protected MavenDependency[] rootDependencies() {
065        return new MavenDependency[] { getRunnerDependency(), getComponentRuntimeBeamDependency() };
066    }
067
068    protected MavenDependency getRunnerDependency() {
069        return createDependency(rootDependencyBase() + ":jar:" + beamVersion, RUNTIME, false);
070    }
071
072    protected MavenDependency getComponentRuntimeBeamDependency() {
073        return createDependency("org.talend.sdk.component:component-runtime-beam:jar:" + kitVersion, RUNTIME, false);
074    }
075
076    @Override
077    public String getName() {
078        return super.getName().replace("Runner", "");
079    }
080
081    protected abstract String rootDependencyBase();
082}