mirror of
https://github.com/facebookresearch/ReAgent.git
synced 2026-06-16 12:44:41 +00:00
Pradeep Kumar Srinivasan
15ec3d92f1
Reviewed By: grievejia Differential Revision: D24655770 fbshipit-source-id: 1fec713e37a8edc4ef0a9855f55c52bcf28eb275
84 lines
2.5 KiB
Python
84 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
|
|
import logging
|
|
import os
|
|
import random
|
|
import shutil
|
|
|
|
import numpy as np
|
|
import torch
|
|
from pyspark import SparkConf
|
|
|
|
# pyre-fixme[21]: Could not find module `reagent.workflow.spark_utils`.
|
|
from reagent.workflow.spark_utils import DEFAULT_SPARK_CONFIG
|
|
|
|
# pyre-fixme[21]: Could not find `sparktestingbase`.
|
|
from sparktestingbase.sqltestcase import SQLTestCase
|
|
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(level=logging.INFO)
|
|
|
|
# path to local hive metastore
|
|
HIVE_METASTORE = "metastore_db"
|
|
|
|
# for setting seeds
|
|
GLOBAL_TEST_CLASS_COUNTER = 0
|
|
|
|
|
|
# pyre-fixme[11]: Annotation `SQLTestCase` is not defined as a type.
|
|
class ReagentSQLTestBase(SQLTestCase):
|
|
def getConf(self):
|
|
conf = SparkConf()
|
|
for k, v in DEFAULT_SPARK_CONFIG.items():
|
|
conf.set(k, v)
|
|
return conf
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super().setUpClass()
|
|
|
|
# set up the seed for the class to prevent
|
|
# clashing random table names for example
|
|
global GLOBAL_TEST_CLASS_COUNTER
|
|
cls.test_class_seed = GLOBAL_TEST_CLASS_COUNTER
|
|
logger.info(f"Allocating seed {cls.test_class_seed} to {cls.__name__}.")
|
|
GLOBAL_TEST_CLASS_COUNTER += 1
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
assert not os.path.isdir(
|
|
HIVE_METASTORE
|
|
), f"{HIVE_METASTORE} already exists! Try deleting it."
|
|
|
|
random.seed(self.test_class_seed)
|
|
torch.manual_seed(self.test_class_seed)
|
|
np.random.seed(self.test_class_seed)
|
|
logging.basicConfig()
|
|
|
|
def assertEq(self, series_a, arr_b):
|
|
""" Assert panda series is equal to np array """
|
|
arr_a = np.array(series_a.tolist())
|
|
np.testing.assert_equal(arr_a, arr_b)
|
|
|
|
def assertAllClose(self, series_a, arr_b):
|
|
""" Assert panda series is allclose to np array """
|
|
arr_a = np.array(series_a.tolist())
|
|
np.testing.assert_allclose(arr_a, arr_b)
|
|
|
|
def assertEqWithPresence(self, series_a, presence, arr_b):
|
|
""" Assert panda series given presence array is equal to np array """
|
|
arr_a = np.array(series_a.tolist())
|
|
present_a = arr_a[presence]
|
|
present_b = arr_b[presence]
|
|
np.testing.assert_equal(present_a, present_b)
|
|
|
|
def tearDown(self):
|
|
super().tearDown()
|
|
|
|
# removes Derby from last runs
|
|
if os.path.isdir(HIVE_METASTORE):
|
|
shutil.rmtree(HIVE_METASTORE)
|