<commit_before>import * as ts from '@terascope/utils';
import shortid from 'shortid';
import * as i from './interfaces';
import { CoreContext } from './core-context';
import validateConfigs from './validate-configs';

interface ClientFactoryFns {
    [prop: string]: i.ClientFactoryFn;
}

export interface CachedClients {
    [prop: string]: any;
}

export interface TestClientConfig {
    type: string;
    create: i.ClientFactoryFn;
    config?: object;
    endpoint?: string;
}

export interface TestClientsByEndpoint {
    [endpoint: string]: any;
}

export interface TestClients {
    [type: string]: TestClientsByEndpoint;
}

export interface TestContextAPIs {
    setTestClients(clients: TestClientConfig[]): void;
    getTestClients(): TestClients;
}

type GetKeyOpts = {
    type: string;
    endpoint?: string;
};

function getKey(opts: GetKeyOpts) {
    const { type, endpoint = 'default' } = opts;
    if (!ts.isString(type)) throw new Error('A type must be specified when registering a Client');
    return `${type}:${endpoint}`;
}

function setConnectorConfig<T extends Record<string, any>>(
    sysconfig: i.FoundationSysConfig<{}>,
    opts: GetKeyOpts,
    config: T,
    override = true
): T {
    const { type, endpoint = 'default' } = opts;
    const { connectors } = sysconfig.terafoundation;
    if (connectors[type] == null) connectors[type] = {};
    if (connectors[type][endpoint] == null) {
        connectors[type][endpoint] = config;
    } else if (override) {
        connectors[type][endpoint] = config;
    }
    return connectors[type][endpoint];
}

const _cachedClients = new WeakMap<TestContext, CachedClients>();
const _createClientFns = new WeakMap<TestContext, ClientFactoryFns>();

export interface TestContextOptions<S> {
    name?: string;
    assignment?: any;
    clients?: TestClientConfig[];
    sysconfig?: S;
}

export class TestContext<
    S = {},
    A = {},
    D extends string = string,
> extends CoreContext<S, A & TestContextAPIs, D> {
    constructor(options: TestContextOptions<S> = {}) {
        const sysconfig: ts.PartialDeep<i.FoundationSysConfig<S>> = {
            terafoundation: {
                connectors: {
                    elasticsearch: {
                        default: {},
                    },
                },
            },
            ...options.sysconfig
        };

        const config: i.FoundationConfig<S, A & TestContextAPIs, D> = {
            name: options.name || 'test-context',
        };
        const cluster: i.Cluster = {
            isMaster: false,
            worker: {
                id: shortid.generate(),
            }
        } as any;

        super(config, cluster, validateConfigs(cluster, config, sysconfig));

        // eslint-disable-next-line @typescript-eslint/no-this-alias
        const ctx = this;
        _cachedClients.set(this, {});<commit_msg>more flexible interfaces for TestContext<commit_after>import * as ts from '@terascope/utils';
import shortid from 'shortid';
import * as i from './interfaces';
import { CoreContext } from './core-context';
import validateConfigs from './validate-configs';

interface ClientFactoryFns {
    [prop: string]: i.ClientFactoryFn;
}

export interface CachedClients {
    [prop: string]: any;
}

export interface TestClientConfig {
    type: string;
    create: i.ClientFactoryFn;
    config?: object;
    endpoint?: string;
}

export interface TestClientsByEndpoint {
    [endpoint: string]: any;
}

export interface TestClients {
    [type: string]: TestClientsByEndpoint;
}

export interface TestContextAPIs {
    setTestClients(clients: TestClientConfig[]): void;
    getTestClients(): TestClients;
}

type GetKeyOpts = {
    type: string;
    endpoint?: string;
};

function getKey(opts: GetKeyOpts) {
    const { type, endpoint = 'default' } = opts;
    if (!ts.isString(type)) throw new Error('A type must be specified when registering a Client');
    return `${type}:${endpoint}`;
}

function setConnectorConfig<T extends Record<string, any>>(
    sysconfig: i.FoundationSysConfig<{}>,
    opts: GetKeyOpts,
    config: T,
    override = true
): T {
    const { type, endpoint = 'default' } = opts;
    const { connectors } = sysconfig.terafoundation;
    if (connectors[type] == null) connectors[type] = {};
    if (connectors[type][endpoint] == null) {
        connectors[type][endpoint] = config;
    } else if (override) {
        connectors[type][endpoint] = config;
    }
    return connectors[type][endpoint];
}

const _cachedClients = new WeakMap<TestContext, CachedClients>();
const _createClientFns = new WeakMap<TestContext, ClientFactoryFns>();

export interface TestContextOptions<S> {
    name?: string;
    assignment?: any;
    clients?: TestClientConfig[];
    sysconfig?: ts.PartialDeep<i.FoundationSysConfig<S>>;
}

export class TestContext<
    S = {},
    A = {},
    D extends string = string,
> extends CoreContext<S, A & TestContextAPIs, D> {
    constructor(options: TestContextOptions<S> = {}) {
        const sysconfig: ts.PartialDeep<i.FoundationSysConfig<S>> = {
            terafoundation: {
                connectors: {
                    elasticsearch: {
                        default: {},
                    },
                },
                ...ts.get(options.sysconfig, 'terafoundation', {})
            },
            ...options.sysconfig
        };

        const config: i.FoundationConfig<S, A & TestContextAPIs, D> = {
            name: options.name || 'test-context',
        };
        const cluster: i.Cluster = {
            isMaster: false,
            worker: {
                id: shortid.generate(),
            }
        } as any;

        super(config, cluster, validateConfigs(cluster, config, sysconfig));

        // eslint-disable-next-line @typescript-eslint/no-this-alias
        const ctx = this;
        _cachedClients.set(this, {});
{'fraction_non_alphanumeric': 0.12629964206579172, 'fraction_numerical': 0.0, 'mean_word_length': 2.749520766773163, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6709281', 'n_tokens_mistral': 1709, 'n_tokens_neox': 1697, 'n_words': 418}