1515#
1616import pytest
1717import uuid
18+ import os
19+ import datetime
1820
1921TEST_FAMILY = "test-family"
2022TEST_FAMILY_2 = "test-family-2"
2123TEST_AGGREGATE_FAMILY = "test-aggregate-family"
2224
2325
26+ # authorized view subset to allow all qualifiers
27+ ALLOW_ALL = ""
28+ ALL_QUALIFIERS = {"qualifier_prefixes" : [ALLOW_ALL ]}
29+
30+
2431class SystemTestRunner :
2532 """
2633 configures a system test class with configuration for clusters/tables/etc
2734
2835 used by standard system tests, and metrics tests
2936 """
3037
31- @pytest .fixture (scope = "session" )
32- def init_table_id (self ):
33- """
34- The table_id to use when creating a new test table
35- """
36- return f"test-table-{ uuid .uuid4 ().hex }
37-
3838 @pytest .fixture (scope = "session" )
3939 def cluster_config (self , project_id ):
4040 """
@@ -68,3 +68,132 @@ def column_family_config(self):
6868 value_type = types .Type (aggregate_type = int_aggregate_type )
6969 ),
7070 }
71+
72+ @pytest .fixture (scope = "session" )
73+ def start_timestamp (self ):
74+ """
75+ A timestamp taken before any tests are run. Used to fetch back metrics relevant to the tests
76+ """
77+ return datetime .datetime .now (datetime .timezone .utc )
78+
79+ def _generate_table_id (self ):
80+ return f"test-table-{ uuid .uuid4 ().hex }
81+
82+ @pytest .fixture (scope = "session" )
83+ def table_id (
84+ self ,
85+ admin_client ,
86+ project_id ,
87+ instance_id ,
88+ column_family_config ,
89+ init_table_id ,
90+ column_split_config ,
91+ start_timestamp ,
92+ ):
93+ """
94+ Returns BIGTABLE_TEST_TABLE if set, otherwise creates a new temporary table for the test session
95+
96+ Args:
97+ - admin_client: Client for interacting with the Table Admin API. Supplied by the admin_client fixture.
98+ - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture.
99+ - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture.
100+ - init_column_families: A list of column families to initialize the table with, if pre-initialized table is not given with BIGTABLE_TEST_TABLE.
101+ Supplied by the init_column_families fixture.
102+ - init_table_id: The table ID to give to the test table, if pre-initialized table is not given with BIGTABLE_TEST_TABLE.
103+ Supplied by the init_table_id fixture.
104+ - column_split_config: A list of row keys to use as initial splits when creating the test table.
105+ - start_timestamp: accessed when building table to ensure timestamp data is loaded before tests are run
106+ """
107+ from google .api_core import exceptions
108+ from google .api_core import retry
109+
110+ # use user-specified instance if available
111+ user_specified_table = os .getenv ("BIGTABLE_TEST_TABLE" )
112+ if user_specified_table :
113+ print ("Using user-specified table: {}" .format (user_specified_table ))
114+ yield user_specified_table
115+ return
116+
117+ retry = retry .Retry (
118+ predicate = retry .if_exception_type (exceptions .FailedPrecondition )
119+ )
120+ try :
121+ parent_path = f"projects/{ project_id } { instance_id }
122+ print (f"Creating table: { parent_path } { init_table_id } )
123+ admin_client .table_admin_client .create_table (
124+ request = {
125+ "parent" : parent_path ,
126+ "table_id" : init_table_id ,
127+ "table" : {"column_families" : column_family_config },
128+ "initial_splits" : [{"key" : key } for key in column_split_config ],
129+ },
130+ retry = retry ,
131+ )
132+ except exceptions .AlreadyExists :
133+ pass
134+ yield init_table_id
135+ print (f"Deleting table: { parent_path } { init_table_id } )
136+ try :
137+ admin_client .table_admin_client .delete_table (
138+ name = f"{ parent_path } { init_table_id }
139+ )
140+ except exceptions .NotFound :
141+ print (f"Table { init_table_id } )
142+
143+ @pytest .fixture (scope = "session" )
144+ def authorized_view_id (
145+ self ,
146+ admin_client ,
147+ project_id ,
148+ instance_id ,
149+ table_id ,
150+ ):
151+ """
152+ Creates and returns a new temporary authorized view for the test session
153+
154+ Args:
155+ - admin_client: Client for interacting with the Table Admin API. Supplied by the admin_client fixture.
156+ - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture.
157+ - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture.
158+ - table_id: The ID of the table to create the authorized view for. Supplied by the table_id fixture.
159+ """
160+ from google .api_core import exceptions
161+ from google .api_core import retry
162+
163+ retry = retry .Retry (
164+ predicate = retry .if_exception_type (exceptions .FailedPrecondition )
165+ )
166+ new_view_id = uuid .uuid4 ().hex [:8 ]
167+ parent_path = f"projects/{ project_id } { instance_id } { table_id }
168+ new_path = f"{ parent_path } { new_view_id }
169+ try :
170+ print (f"Creating view: { new_path } )
171+ admin_client .table_admin_client .create_authorized_view (
172+ request = {
173+ "parent" : parent_path ,
174+ "authorized_view_id" : new_view_id ,
175+ "authorized_view" : {
176+ "subset_view" : {
177+ "row_prefixes" : [ALLOW_ALL ],
178+ "family_subsets" : {
179+ TEST_FAMILY : ALL_QUALIFIERS ,
180+ TEST_FAMILY_2 : ALL_QUALIFIERS ,
181+ TEST_AGGREGATE_FAMILY : ALL_QUALIFIERS ,
182+ },
183+ },
184+ },
185+ },
186+ retry = retry ,
187+ )
188+ except exceptions .AlreadyExists :
189+ pass
190+ except exceptions .MethodNotImplemented :
191+ # will occur when run in emulator. Pass empty id
192+ new_view_id = None
193+ yield new_view_id
194+ if new_view_id :
195+ print (f"Deleting view: { new_path } )
196+ try :
197+ admin_client .table_admin_client .delete_authorized_view (name = new_path )
198+ except exceptions .NotFound :
199+ print (f"View { new_view_id } )
0 commit comments