def test_query_sessions(self):
session = self.dal.session.create(Session(self.session_input_dict))
assert len(self.dal.session.query({"id": session.id})) == 1
_ = self.dal.session.create(Session(self.session_input_dict))
assert len(
self.dal.session.query({
"name": self.session_input_dict['name']
})) == 2
def test_query_sessions_range_query(self):
_ = self.dal.session.create(Session(self.session_input_dict))
_ = self.dal.session.create(Session(self.session_input_dict))
_ = self.dal.session.create(Session(self.session_input_dict))
sessions = self.dal.session.query(
{}, sort_key="created_at", sort_order="descending")
result = self.dal.session.query({
"created_at": {
"$lt": sessions[1].created_at.strftime('%Y-%m-%dT%H:%M:%S.%fZ')
}
})
assert len(sessions) == 3
assert len(result) == 1
def test_query_sessions_current(self):
session_1 = self.dal.session.create(Session(self.session_input_dict))
session_2 = self.dal.session.create(
Session({
"name": "test",
"model_id": "test",
"current": False
}))
current_sessions = self.dal.session.query({"current": True})
noncurrent_sessions = self.dal.session.query({"current": False})
assert len(current_sessions) == 1
assert current_sessions[0] == session_1
assert len(noncurrent_sessions) == 1
assert noncurrent_sessions[0] == session_2
def setup_method(self):
# provide mountable tmp directory for docker
tempfile.tempdir = "/tmp" if not platform.system(
) == "Windows" else None
test_datmo_dir = os.environ.get('TEST_DATMO_DIR',
tempfile.gettempdir())
self.temp_dir = tempfile.mkdtemp(dir=test_datmo_dir)
self.datadriver = BlitzDBDALDriver("file", self.temp_dir)
self.dal = LocalDAL(self.datadriver)
model_name = "model_1"
model = self.dal.model.create(Model({"name": model_name}))
session_name = "session_1"
session = self.dal.session.create(
Session({
"name": session_name,
"model_id": model.id
}))
self.task_input_dict = {
"model_id": model.id,
"session_id": session.id,
"command": "task_1",
"start_time": datetime.utcnow(),
"end_time": datetime.utcnow(),
"duration": 0.004,
}
def test_init_with_id(self):
session_entity = Session(self.input_dict)
for k, v in self.input_dict.items():
assert getattr(session_entity, k) == v
assert session_entity.created_at
assert session_entity.updated_at
def test_current_session(self):
# Test failure case
failed = False
try:
_ = self.base_controller.current_session
except DatmoModelNotInitialized:
failed = True
assert failed
# Test success case
self.base_controller.dal.model.create(
Model({
"name": "test",
"description": "test"
}))
_ = self.base_controller.model
self.base_controller.dal.session.create(
Session({
"name": "test",
"model_id": "test",
"current": True
}))
session = self.base_controller.current_session
assert session.id
assert session.name == "test"
assert session.model_id == "test"
assert session.current == True
def setup_method(self):
# provide mountable tmp directory for docker
tempfile.tempdir = "/tmp" if not platform.system(
) == "Windows" else None
test_datmo_dir = os.environ.get('TEST_DATMO_DIR',
tempfile.gettempdir())
self.temp_dir = tempfile.mkdtemp(dir=test_datmo_dir)
self.datadriver = BlitzDBDALDriver("file", self.temp_dir)
self.dal = LocalDAL(self.datadriver)
model_name = "model_1"
model = self.dal.model.create(Model({"name": model_name}))
session_name = "session_1"
session = self.dal.session.create(
Session({
"name": session_name,
"model_id": model.id
}))
self.snapshot_input_dict = {
"model_id": model.id,
"session_id": session.id,
"message": "my message",
"code_id": "code_id",
"environment_id": "environment_id",
"file_collection_id": "file_collection_id",
"config": {
"test": 0.45
},
"stats": {
"test": 0.98
}
}
def test_create_session_by_dictionary(self):
session = self.dal.session.create(Session(self.session_input_dict))
assert session.id
assert session.name == self.session_input_dict['name']
assert session.created_at
assert session.updated_at
session_2 = self.dal.session.create(Session(self.session_input_dict))
assert session_2.id != session.id
test_session_input_dict = self.session_input_dict.copy()
test_session_input_dict['id'] = "session_id"
session_3 = self.dal.session.create(Session(test_session_input_dict))
assert session_3.id == test_session_input_dict['id']
def test_delete_session(self):
session = self.dal.session.create(Session(self.session_input_dict))
self.dal.session.delete(session.id)
deleted = False
try:
self.dal.session.get_by_id(session.id)
except EntityNotFound:
deleted = True
assert deleted
def test_get_by_id_session_new_driver_instance(self):
session = self.dal.session.create(Session(self.session_input_dict))
# create new dal with new driver instance (fails)
new_driver_instance = BlitzDBDALDriver("file", self.temp_dir)
new_dal_instance = LocalDAL(new_driver_instance)
new_session_1 = new_dal_instance.session.get_by_id(session.id)
assert new_session_1.id == session.id
# create new dal instance with same driver (success)
new_dal_instance = LocalDAL(self.datadriver)
new_session_2 = new_dal_instance.session.get_by_id(session.id)
assert new_session_2.id == session.id
def test_create_update_key_to_same_value(self):
session_1 = self.dal.session.create(Session(self.session_input_dict))
session_2 = self.dal.session.create(
Session({
"name": "test",
"model_id": "test",
"current": False
}))
next_session = self.dal.session.query({"model_id": session_2.model_id})
next_session = next_session[0]
# get current session and update to false if present
current_session = self.dal.session.query({
"model_id": session_1.model_id
})
if len(current_session) == 1:
current_session = current_session[0]
current_session.model_id = session_2.model_id
self.dal.session.update(current_session)
# update the next session to the current one
next_session.model_id = session_1.model_id
self.dal.session.update(next_session)
# ensure we still get the right result
current_sessions = self.dal.session.query({
"model_id": session_1.model_id
})
noncurrent_sessions = self.dal.session.query({
"model_id": session_2.model_id
})
assert len(current_sessions) == 1
assert current_sessions[0] == session_2
assert len(noncurrent_sessions) == 1
assert noncurrent_sessions[0] == session_1
def test_update_session(self):
session = self.dal.session.create(Session(self.session_input_dict))
# Update required and optional parameters
updated_session_input_dict = self.session_input_dict.copy()
updated_session_input_dict['id'] = session.id
updated_session_input_dict['name'] = "session_yo"
updated_session_input_dict['created_at'] = datetime.utcnow()
updated_session = self.dal.session.update(updated_session_input_dict)
assert session.id == updated_session.id
assert session.updated_at < updated_session.updated_at
assert updated_session.name == updated_session_input_dict['name']
assert updated_session.created_at == updated_session_input_dict[
'created_at']
def init(self, name, description):
# Error if name is not given
if not name:
raise RequiredArgumentMissing(
__("error", "controller.project.init.arg", "name"))
# Create the Model, is it new or update?
is_new_model = False
if not self.model:
_ = self.dal.model.create(
Model({
"name": name,
"description": description
}))
is_new_model = True
else:
self._model = self.dal.model.update({
"id": self.model.id,
"name": name,
"description": description
})
# Initialize Code Manager if needed
if not self.code_driver.is_initialized:
self.code_driver.init()
# Initialize File Manager if needed
if not self.file_driver.is_initialized:
self.file_driver.init()
# Initialize Environment Manager if needed
if not self.environment_driver.is_initialized:
self.environment_driver.init()
# Build the initial default Environment (NOT NECESSARY)
# self.environment_driver.build_image(tag="datmo-" + \
# self.model.name)
# Add in Project template files if specified
# TODO: Add in project template files
# Create and set current session
if is_new_model:
# Create new default session
_ = self.dal.session.create(
Session({
"name": "default",
"model_id": self.model.id,
"current": True
}))
else:
if not self.current_session:
default_session_obj = self.dal.session.query({
"name":
"default",
"model_id":
self.model.id
})
if not default_session_obj:
raise SessionDoesNotExistException(
__("error", "controller.project.init"))
# Update default session to be current
self.dal.session.update({
"id": default_session_obj.id,
"current": True
})
return True
def init(self, name, description):
""" Initialize the project
This function will initialize the project or reinitialize it the project is
already initialized.
Parameters
----------
name : str
description : str
Returns
-------
bool
Raises
------
SessionDoesNotExist
"""
# Create the Model, is it new or update?
is_new_model = False
old_model = self.model
if not self.model:
is_new_model = True
try:
# Always validate inputs to the init function
validate("create_project", {
"name": name,
"description": description
})
# Create model if new else update
if is_new_model:
_ = self.dal.model.create(
Model({
"name": name,
"description": description
}))
else:
self._model = self.dal.model.update({
"id": self.model.id,
"name": name,
"description": description
})
# Initialize Code Driver if needed
if not self.code_driver.is_initialized:
self.code_driver.init()
# Initialize File Driver if needed
if not self.file_driver.is_initialized:
self.file_driver.init()
# Initialize Environment Driver if needed
# (not required but will warn if not present)
try:
if not self.environment_driver.is_initialized:
self.environment_driver.init()
except EnvironmentInitFailed:
self.logger.warning(
__("warn", "controller.general.environment.failed"))
# Build the initial default Environment (NOT NECESSARY)
# self.environment_driver.build_image(tag="datmo-" + \
# self.model.name)
# Create and set current session
if is_new_model:
# Create new default session
_ = self.dal.session.create(
Session({
"name": "default",
"model_id": self.model.id,
"current": True
}))
else:
if not self.current_session:
default_session_objs = self.dal.session.query({
"name": "default",
"model_id": self.model.id
})
if not default_session_objs:
# Creating a default session since none exists
_ = self.dal.session.create(
Session({
"name": "default",
"model_id": self.model.id,
"current": True
}))
else:
# Update default session to be current
default_session_obj = default_session_objs[0]
self.dal.session.update({
"id": default_session_obj.id,
"current": True
})
return True
except Exception:
# if any error occurred with new model, ensure no initialize occurs and raise previous error
# if any error occurred with existing model, ensure no updates were made, raise previous error
if is_new_model:
self.cleanup()
else:
self._model = self.dal.model.update({
"id": old_model.id,
"name": old_model.name,
"description": old_model.description
})
raise
def test_sort_sessions(self):
session_1 = self.dal.session.create(Session(self.session_input_dict))
session_2 = self.dal.session.create(Session(self.session_input_dict))
# Sorting of snapshot in descending
items = self.dal.session.query(
{
"name": self.session_input_dict['name']
},
sort_key="created_at",
sort_order="descending")
assert items[0].created_at == session_2.created_at
# Sorting of snapshot in ascending
items = self.dal.session.query(
{
"name": self.session_input_dict["name"]
},
sort_key="created_at",
sort_order="ascending")
assert items[0].created_at == session_1.created_at
# Wrong order being passed in
failed = False
try:
_ = self.dal.session.query(
{
"name": self.session_input_dict['name']
},
sort_key="created_at",
sort_order="wrong_order")
except InvalidArgumentType:
failed = True
assert failed
# Wrong key and order being passed in
failed = False
try:
_ = self.dal.session.query(
{
"name": self.session_input_dict["name"]
},
sort_key="wrong_key",
sort_order="wrong_order")
except InvalidArgumentType:
failed = True
assert failed
# wrong key and right order being passed in
expected_items = self.dal.session.query(
{
"name": self.session_input_dict["name"]
},
sort_key="created_at",
sort_order="ascending")
items = self.dal.session.query(
{
"name": self.session_input_dict["name"]
},
sort_key="wrong_key",
sort_order="ascending")
expected_ids = [item.id for item in expected_items]
ids = [item.id for item in items]
assert set(expected_ids) == set(ids)
def test_to_dictionary(self):
session_entity = Session(self.input_dict)
output_dict = session_entity.to_dictionary()
for k, v in output_dict.items():
assert v == getattr(session_entity, k)
def test_eq(self):
session_entity_1 = Session(self.input_dict)
session_entity_2 = Session(self.input_dict)
assert session_entity_1 == session_entity_2
def test_get_by_shortened_id_session(self):
session = self.dal.session.create(Session(self.session_input_dict))
result = self.dal.session.get_by_shortened_id(session.id[:10])
assert session.id == result.id
