def test_overwritten_values(self):
model_io_stub = stub_io.ModelIoStub()
get_models = model_io_stub.get_models_ptransform
options = job_options.JobOptions(model_getter=get_models)
self.assertIs(options.model_getter, get_models)
def test_get_models_ptransform(self):
stub = stub_io.ModelIoStub()
test_models = [
base_models.BaseModel(
id='a', created_on=self.NOW, last_updated=self.NOW),
base_models.BaseModel(
id='b', created_on=self.NOW, last_updated=self.NOW),
base_models.BaseModel(
id='c', created_on=self.NOW, last_updated=self.NOW),
]
stub.put_multi(test_models)
self.assert_pcoll_equal(
self.pipeline | stub.get_models_ptransform(), test_models)
def test_get_models_returns_nothing_when_stub_is_empty(self):
stub = stub_io.ModelIoStub()
self.assert_pcoll_empty(self.pipeline | stub.get_models())
def __init__(self, *args, **kwargs):
super(JobTestBase, self).__init__(*args, **kwargs)
self.model_io_stub = stub_io.ModelIoStub()
self.pipeline.options.view_as(job_options.JobOptions).model_getter = (
self.model_io_stub.get_models_ptransform)