mirror of
https://github.com/neondatabase/neon.git
synced 2026-01-07 05:22:56 +00:00
30a7dd630c
## Problem `TYPE_CHECKING` is used inconsistently across Python tests. ## Summary of changes - Update `ruff`: 0.7.0 -> 0.11.2 - Enable TC (flake8-type-checking): https://docs.astral.sh/ruff/rules/#flake8-type-checking-tc - (auto)fix all new issues
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import psycopg2
|
|
import pytest
|
|
|
|
if TYPE_CHECKING:
|
|
from fixtures.neon_fixtures import (
|
|
NeonProxy,
|
|
VanillaPostgres,
|
|
)
|
|
|
|
TABLE_NAME = "neon_control_plane.endpoints"
|
|
|
|
|
|
# Proxy uses the same logic for psql and websockets.
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_psql_allowed_ips(static_proxy: NeonProxy, vanilla_pg: VanillaPostgres):
|
|
# Shouldn't be able to connect to this project
|
|
vanilla_pg.safe_psql(
|
|
f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('private-project', '8.8.8.8')"
|
|
)
|
|
# Should be able to connect to this project
|
|
vanilla_pg.safe_psql(
|
|
f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('generic-project', '::1,127.0.0.1')"
|
|
)
|
|
|
|
def check_cannot_connect(**kwargs):
|
|
with pytest.raises(psycopg2.Error) as exprinfo:
|
|
static_proxy.safe_psql(**kwargs)
|
|
text = str(exprinfo.value).strip()
|
|
assert "not allowed to connect" in text
|
|
|
|
# no SNI, deprecated `options=project` syntax (before we had several endpoint in project)
|
|
check_cannot_connect(query="select 1", sslsni=0, options="project=private-project")
|
|
|
|
# no SNI, new `options=endpoint` syntax
|
|
check_cannot_connect(query="select 1", sslsni=0, options="endpoint=private-project")
|
|
|
|
# with SNI
|
|
check_cannot_connect(query="select 1", host="private-project.local.neon.build")
|
|
|
|
# no SNI, deprecated `options=project` syntax (before we had several endpoint in project)
|
|
out = static_proxy.safe_psql(query="select 1", sslsni=0, options="project=generic-project")
|
|
assert out[0][0] == 1
|
|
|
|
# no SNI, new `options=endpoint` syntax
|
|
out = static_proxy.safe_psql(query="select 1", sslsni=0, options="endpoint=generic-project")
|
|
assert out[0][0] == 1
|
|
|
|
# with SNI
|
|
out = static_proxy.safe_psql(query="select 1", host="generic-project.local.neon.build")
|
|
assert out[0][0] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_http_allowed_ips(static_proxy: NeonProxy, vanilla_pg: VanillaPostgres):
|
|
static_proxy.safe_psql("create user http_auth with password 'http' superuser")
|
|
|
|
# Shouldn't be able to connect to this project
|
|
vanilla_pg.safe_psql(
|
|
f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('proxy', '8.8.8.8')"
|
|
)
|
|
|
|
def query(status: int, query: str, *args):
|
|
static_proxy.http_query(
|
|
query,
|
|
args,
|
|
user="http_auth",
|
|
password="http",
|
|
expected_code=status,
|
|
)
|
|
|
|
query(400, "select 1;") # ip address is not allowed
|
|
# Should be able to connect to this project
|
|
vanilla_pg.safe_psql(
|
|
f"UPDATE {TABLE_NAME} SET allowed_ips = '8.8.8.8,127.0.0.1' WHERE endpoint_id = 'proxy'"
|
|
)
|
|
query(200, "select 1;") # should work now
|