Repository: apache/airflow
File: providers/edge3/tests/unit/edge3/cli/test_worker.py (Python)

Code before:
        mock_push_log_chunks.assert_called_once()
        assert len(worker_with_job.jobs) == 1  # no new job added (was removed at the end...)
        mock_logs_push.assert_called_once()

    @time_machine.travel(datetime.now(), tick=False)
    @patch("airflow.providers.edge3.cli.worker.logs_push")
    @pytest.mark.asyncio
    async def test_push_logs_in_chunks(self, mock_logs_push, worker_with_job: EdgeWorker):
        job = EdgeWorker.jobs[0]
        await anyio.Path(job.logfile).write_text("some log content")
        with conf_vars({("edge", "api_url"): "https://invalid-api-test-endpoint"}):
            await worker_with_job._push_logs_in_chunks(job)

        assert len(EdgeWorker.jobs) == 1
        mock_logs_push.assert_called_once_with(
            task=job.edge_job.key, log_chunk_time=timezone.utcnow(), log_chunk_data="some log content"
        )

    @time_machine.travel(datetime.now(), tick=False)
    @patch("airflow.providers.edge3.cli.worker.logs_push")
    @pytest.mark.asyncio
    async def test_check_running_jobs_log_push_increment(self, mock_logs_push, worker_with_job: EdgeWorker):
        job = EdgeWorker.jobs[0]
        await anyio.Path(job.logfile).write_text("hello ")
        job.logsize = (await anyio.Path(job.logfile).stat()).st_size
        await anyio.Path(job.logfile).write_text("hello world")
        with conf_vars({("edge", "api_url"): "https://invalid-api-test-endpoint"}):
            await worker_with_job._push_logs_in_chunks(job)
        assert len(EdgeWorker.jobs) == 1
        mock_logs_push.assert_called_once_with(
            task=job.edge_job.key, log_chunk_time=timezone.utcnow(), log_chunk_data="world"
        )

    @time_machine.travel(datetime.now(), tick=False)
    @patch("airflow.providers.edge3.cli.worker.logs_push")
    @patch.object(worker_module, "push_log_chunk_size", 4)
    @pytest.mark.asyncio
    async def test_check_running_jobs_log_push_chunks(self, mock_logs_push, worker_with_job: EdgeWorker):
        job = EdgeWorker.jobs[0]
        job.logfile.write_bytes("log1log2ülog3".encode("latin-1"))
        with conf_vars({("edge", "api_url"): "https://invalid-api-test-endpoint"}):
            await worker_with_job._push_logs_in_chunks(job)
        assert len(EdgeWorker.jobs) == 1
        calls = mock_logs_push.call_args_list
        assert len(calls) == 4
        assert calls[0] == call(
            task=job.edge_job.key, log_chunk_time=timezone.utcnow(), log_chunk_data="log1"
        )
        assert calls[1] == call(
            task=job.edge_job.key, log_chunk_time=timezone.utcnow(), log_chunk_data="log2"
        )

Review comment (suggestion):
the test creates anyio.Path(job.logfile) three times in test_check_running_jobs_log_push_increment instead of caching it like the production code does.
```suggestion
        aio_logfile = anyio.Path(job.logfile)
        await aio_logfile.write_text("hello ")
        job.logsize = (await aio_logfile.stat()).st_size
        await aio_logfile.write_text("hello world")
```

Code after:
        mock_push_log_chunks.assert_called_once()
        assert len(worker_with_job.jobs) == 1  # no new job added (was removed at the end...)
        mock_logs_push.assert_called_once()

    @time_machine.travel(datetime.now(), tick=False)
    @patch("airflow.providers.edge3.cli.worker.logs_push")
    @pytest.mark.asyncio
    async def test_push_logs_in_chunks(self, mock_logs_push, worker_with_job: EdgeWorker):
        job = EdgeWorker.jobs[0]
        await anyio.Path(job.logfile).write_text("some log content")
        with conf_vars({("edge", "api_url"): "https://invalid-api-test-endpoint"}):
            await worker_with_job._push_logs_in_chunks(job)

        assert len(EdgeWorker.jobs) == 1
        mock_logs_push.assert_called_once_with(
            task=job.edge_job.key, log_chunk_time=timezone.utcnow(), log_chunk_data="some log content"
        )

    @time_machine.travel(datetime.now(), tick=False)
    @patch("airflow.providers.edge3.cli.worker.logs_push")
    @pytest.mark.asyncio
    async def test_check_running_jobs_log_push_increment(self, mock_logs_push, worker_with_job: EdgeWorker):
        job = EdgeWorker.jobs[0]
        aio_logfile = anyio.Path(job.logfile)
        await aio_logfile.write_text("hello ")
        job.logsize = (await aio_logfile.stat()).st_size
        await aio_logfile.write_text("hello world")
        with conf_vars({("edge", "api_url"): "https://invalid-api-test-endpoint"}):
            await worker_with_job._push_logs_in_chunks(job)
        assert len(EdgeWorker.jobs) == 1
        mock_logs_push.assert_called_once_with(
            task=job.edge_job.key, log_chunk_time=timezone.utcnow(), log_chunk_data="world"
        )

    @time_machine.travel(datetime.now(), tick=False)
    @patch("airflow.providers.edge3.cli.worker.logs_push")
    @patch.object(worker_module, "push_log_chunk_size", 4)
    @pytest.mark.asyncio
    async def test_check_running_jobs_log_push_chunks(self, mock_logs_push, worker_with_job: EdgeWorker):
        job = EdgeWorker.jobs[0]
        job.logfile.write_bytes("log1log2ülog3".encode("latin-1"))
        with conf_vars({("edge", "api_url"): "https://invalid-api-test-endpoint"}):
            await worker_with_job._push_logs_in_chunks(job)
        assert len(EdgeWorker.jobs) == 1
        calls = mock_logs_push.call_args_list
        assert len(calls) == 4
        assert calls[0] == call(
            task=job.edge_job.key, log_chunk_time=timezone.utcnow(), log_chunk_data="log1"
        )
        assert calls[1] == call(
            task=job.edge_job.key, log_chunk_time=timezone.utcnow(), log_chunk_data="log2"