Add missing tests

Author: damklis

airflow/modules/rss_news/rss_news_exporter.py

@@ -5,19 +5,25 @@
 
 class NewsExporter:
  def __init__(self, bootstrap_servers):
- self.producer = KafkaProducer(
- bootstrap_servers=bootstrap_servers,
- value_serializer=lambda x: self._encode_news(x)
+ self._producer = self._connect_producer(
+ bootstrap_servers
  )
 
- def _encode_news(self, value):
- return json.dumps(value).encode("utf-8")
+ def _connect_producer(self, bootstrap_servers):
+ def encode_news(value):
+ return json.dumps(value).encode("utf-8")
+
+ producer = KafkaProducer(
+ bootstrap_servers=bootstrap_servers,
+ value_serializer=lambda x: encode_news(x)
+ )
+ return producer
 
  def __enter__(self):
  return self
 
  def export_news_to_broker(self, topic, record, sleep_time=0.01):
- response = self.producer.send(
+ response = self._producer.send(
  topic,
  value=record
  )
@@ -27,4 +33,4 @@ def export_news_to_broker(self, topic, record, sleep_time=0.01):
  )
 
  def __exit__(self, type, value, traceback):
- self.producer.close()
+ self._producer.close()

airflow/modules/tests/rss_news/test_rss_news_exporter.py

@@ -1,5 +1,14 @@
-from unittest.mock import patch
+from unittest.mock import patch, Mock
 import pytest
+from rss_news import NewsExporter
+
+
+@patch("rss_news.rss_news_exporter.KafkaProducer")
+def test_connect_producer(mock_producer):
+
+ exporter = NewsExporter(["test_broker:9092"])
+
+ assert exporter._producer is not None
 
 
 @patch("rss_news.NewsExporter")
@@ -19,3 +28,24 @@ def test_export_news_to_broker(exporter):
  exporter.export_news_to_broker.assert_called_once_with(
  topic, news
  )
+
+
+@patch("rss_news.rss_news_exporter.KafkaProducer")
+def test_export_news_to_broker_context_manager(mock_producer):
+ mock_producer.send.return_value = Mock()
+ topic = "test_topic"
+ news = {
+ "_id": "test_id",
+ "title": "test_title",
+ "link": "www.test.com",
+ "date": "2020-01-01 00:00:00",
+ "description": "Test",
+ "author": "Test",
+ "language": "pl"
+ }
+
+ with NewsExporter(["test_broker:9092"]) as exporter:
+ exporter.export_news_to_broker(topic, news)
+ exporter._producer.send.assert_called_once_with(
+ topic, value=news
+ )