Skip to content

Commit

Permalink
format
Browse files Browse the repository at this point in the history
  • Loading branch information
elay committed Jun 24, 2024
1 parent 0c1bb14 commit c546d7f
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 11 deletions.
10 changes: 7 additions & 3 deletions pcfuncs/ipban/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,13 @@ def main(mytimer: func.TimerRequest) -> None:
datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat()
)
logger.info("Updating the ip ban list at %s", utc_timestamp)
credential: DefaultAzureCredential = DefaultAzureCredential()
credential: DefaultAzureCredential = DefaultAzureCredential()
with LogsQueryClient(credential) as logs_query_client:
with TableServiceClient(endpoint=settings.storage_account_url, credential=credential) as table_service_client:
with table_service_client.create_table_if_not_exists(settings.banned_ip_table) as table_client:
with TableServiceClient(
endpoint=settings.storage_account_url, credential=credential
) as table_service_client:
with table_service_client.create_table_if_not_exists(
settings.banned_ip_table
) as table_client:
task = UpdateBannedIPTask(logs_query_client, table_client)
task.run()
23 changes: 15 additions & 8 deletions pcfuncs/tests/ipban/test_ipban.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,12 +64,16 @@ def mock_clients(
)

# Use Azurite for unit tests and populate the table with initial data
with TableServiceClient.from_connection_string(CONNECTION_STRING) as table_service_client:
with table_service_client.create_table_if_not_exists(table_name=TEST_BANNED_IP_TABLE) as table_client:

with TableServiceClient.from_connection_string(
CONNECTION_STRING
) as table_service_client:
with table_service_client.create_table_if_not_exists(
table_name=TEST_BANNED_IP_TABLE
) as table_client:

# Pre-populate the banned IP table
populate_banned_ip_table(table_client)

# Yield the clients for use
yield logs_query_client, table_client

Expand All @@ -83,20 +87,23 @@ def integration_clients(
) -> Generator[Tuple[LogsQueryClient, TableClient], Any, None]:
credential: DefaultAzureCredential = DefaultAzureCredential()
with LogsQueryClient(credential) as logs_query_client:
with TableServiceClient(endpoint=settings.storage_account_url, credential=credential) as table_service_client:
with table_service_client.create_table_if_not_exists(TEST_BANNED_IP_TABLE) as table_client:
with TableServiceClient(
endpoint=settings.storage_account_url, credential=credential
) as table_service_client:
with table_service_client.create_table_if_not_exists(
TEST_BANNED_IP_TABLE
) as table_client:

# Pre-populate the banned IP table
populate_banned_ip_table(table_client)

# Yield the clients for use
yield logs_query_client, table_client

# Delete the test table
table_service_client.delete_table(TEST_BANNED_IP_TABLE)



@pytest.mark.integration
def test_update_banned_ip_integration(
integration_clients: Tuple[LogsQueryClient, TableClient]
Expand Down

0 comments on commit c546d7f

Please sign in to comment.