Dmitrii,
I use a different approach, where I tangle the source into files in modules and then I import those modules from other blocks.
This allows me to organize my document with different sections for the code and its tests, which then get exported into their corresponding files.
* Square Function
This function receives a number and returns its square
#+BEGIN_SRC python :tangle ./utils/math.py :mkdirp yes
def square(x):
return x * x
#+END_SRC
** __init__.py (module setup)
#+begin_src python :tangle ./utils/__init__.py :mkdirp yes
from utils.math import square
** Test cases
1. The square of five should be 25
2. The square of zero should be 0
3. The square of a negative number should be positive
#+BEGIN_SRC python :tangle ./utils/test_square.py :mkdirp yes
from utils.math import square
def test_square_of_five():
assert square(5) == 25
def test_square_of_zero():
assert square(0) == 0
def test_square_of_negative():
assert square(-5) > 0
#+END_SRC
*** Run tests
#+begin_src sh :results output raw
pytest ./utils
#+end_src
#+RESULTS:
============================= test session starts ==============================
platform linux -- Python 3.7.3, pytest-4.6.3, py-1.8.0, pluggy-0.12.0
rootdir: /app
collected 3 items
utils/test_square.py ... [100%]
=========================== 3 passed in 0.08 seconds ===========================