def setUp(self):
environ["SUBMARINE_JOB_NAME"] = JOB_NAME
submarine.set_tracking_uri(
"mysql+pymysql://submarine_test:password_test@localhost:3306/submarineDB_test"
)
self.tracking_uri = utils.get_tracking_uri()
self.store = utils.get_sqlalchemy_store(self.tracking_uri)
def __init__(self, tracking_uri=None):
"""
:param tracking_uri: Address of local or remote tracking server. If not provided, defaults
to the service set by ``submarine.tracking.set_tracking_uri``. See
`Where Runs Get Recorded <../tracking.html#where-runs-get-recorded>`_
for more info.
"""
self.tracking_uri = tracking_uri or utils.get_tracking_uri()
self.store = utils.get_sqlalchemy_store(self.tracking_uri)
def test_log_param(tracking_uri_mock):
environ["SUBMARINE_JOB_NAME"] = JOB_NAME
submarine.log_param("name_1", "a", "worker-1")
tracking_uri = utils.get_tracking_uri()
store = utils.get_sqlalchemy_store(tracking_uri)
# Validate params
with store.ManagedSessionMaker() as session:
params = session \
.query(SqlParam) \
.options() \
.filter(SqlParam.job_name == JOB_NAME).all()
assert params[0].key == "name_1"
assert params[0].value == "a"
assert params[0].worker_index == "worker-1"
assert params[0].job_name == JOB_NAME
def test_log_metric(tracking_uri_mock):
environ["SUBMARINE_JOB_NAME"] = JOB_NAME
submarine.log_metric("name_1", 5, "worker-1")
submarine.log_metric("name_1", 6, "worker-2")
tracking_uri = utils.get_tracking_uri()
store = utils.get_sqlalchemy_store(tracking_uri)
# Validate params
with store.ManagedSessionMaker() as session:
metrics = session \
.query(SqlMetric) \
.options() \
.filter(SqlMetric.job_name == JOB_NAME).all()
assert len(metrics) == 2
assert metrics[0].key == "name_1"
assert metrics[0].value == 5
assert metrics[0].worker_index == "worker-1"
assert metrics[0].job_name == JOB_NAME
assert metrics[1].value == 6
assert metrics[1].worker_index == "worker-2"
def test_get_tracking_uri():
env = {
_TRACKING_URI_ENV_VAR: DEFAULT_SUBMARINE_JDBC_URL,
}
with mock.patch.dict(os.environ, env):
assert get_tracking_uri() == DEFAULT_SUBMARINE_JDBC_URL
def setUp(self):
submarine.set_tracking_uri(
"mysql+pymysql://submarine_test:password_test@localhost:3306/submarine_test"
)
self.tracking_uri = utils.get_tracking_uri()
self.store = utils.get_sqlalchemy_store(self.tracking_uri)
