100100from .backup import NodeBackup
101101
102102from .operations .os_ops import ConnectionParams
103+ from .operations .os_ops import OsOperations
103104from .operations .local_ops import LocalOperations
104105from .operations .remote_ops import RemoteOperations
105106
@@ -135,7 +136,7 @@ class PostgresNode(object):
135136 _C_MAX_START_ATEMPTS = 5
136137
137138 def __init__ (self , name = None , base_dir = None , port = None , conn_params : ConnectionParams = ConnectionParams (),
138- bin_dir = None , prefix = None ):
139+ bin_dir = None , prefix = None , os_ops = None ):
139140 """
140141 PostgresNode constructor.
141142
@@ -157,17 +158,20 @@ def __init__(self, name=None, base_dir=None, port=None, conn_params: ConnectionP
157158
158159 # basic
159160 self .name = name or generate_app_name ()
160- if testgres_config .os_ops :
161- self .os_ops = testgres_config .os_ops
162- elif conn_params .ssh_key :
163- self .os_ops = RemoteOperations (conn_params )
161+ if os_ops is None :
162+ os_ops = __class__ ._get_os_ops (conn_params )
164163 else :
165- self .os_ops = LocalOperations (conn_params )
164+ assert conn_params is None
165+ pass
166166
167- self .host = self .os_ops .host
167+ assert os_ops is not None
168+ assert isinstance (os_ops , OsOperations )
169+ self ._os_ops = os_ops
170+
171+ self .host = os_ops .host
168172 self .port = port or utils .reserve_port ()
169173
170- self .ssh_key = self . os_ops .ssh_key
174+ self .ssh_key = os_ops .ssh_key
171175
172176 # defaults for __exit__()
173177 self .cleanup_on_good_exit = testgres_config .node_cleanup_on_good_exit
@@ -204,6 +208,40 @@ def __repr__(self):
204208 return "{}(name='{}', port={}, base_dir='{}')" .format (
205209 self .__class__ .__name__ , self .name , self .port , self .base_dir )
206210
211+ @staticmethod
212+ def _get_os_ops (conn_params : ConnectionParams ) -> OsOperations :
213+ if testgres_config .os_ops :
214+ return testgres_config .os_ops
215+
216+ assert type (conn_params ) == ConnectionParams # noqa: E721
217+
218+ if conn_params .ssh_key :
219+ return RemoteOperations (conn_params )
220+
221+ return LocalOperations (conn_params )
222+
223+ def clone_with_new_name_and_base_dir (self , name : str , base_dir : str ):
224+ assert name is None or type (name ) == str # noqa: E721
225+ assert base_dir is None or type (base_dir ) == str # noqa: E721
226+
227+ assert __class__ == PostgresNode
228+
229+ node = PostgresNode (
230+ name = name ,
231+ base_dir = base_dir ,
232+ conn_params = None ,
233+ bin_dir = self ._bin_dir ,
234+ prefix = self ._prefix ,
235+ os_ops = self ._os_ops )
236+
237+ return node
238+
239+ @property
240+ def os_ops (self ) -> OsOperations :
241+ assert self ._os_ops is not None
242+ assert isinstance (self ._os_ops , OsOperations )
243+ return self ._os_ops
244+
207245 @property
208246 def pid (self ):
209247 """
0 commit comments