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)